import gradio as gr
from huggingface_hub import InferenceClient
import edge_tts
import asyncio
import requests
from bs4 import BeautifulSoup

client = InferenceClient("google/gemma-2b-it")


def search(query):
    term = query
    start = 0
    all_results = []
    max_chars_per_page = 8000
    with requests.Session() as session:
        resp = session.get(
            url="https://www.google.com/search",
            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
            params={"q": term, "num": 3, "udm": 14},
            timeout=5,
            verify=None,
        )
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        result_block = soup.find_all("div", attrs={"class": "g"})
        for result in result_block:
            link = result.find("a", href=True)
            link = link["href"]
            try:
                webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"}, timeout=5, verify=False)
                webpage.raise_for_status()
                visible_text = extract_text_from_webpage(webpage.text)
                if len(visible_text) > max_chars_per_page:
                    visible_text = visible_text[:max_chars_per_page]
                all_results.append({"link": link, "text": visible_text})
            except requests.exceptions.RequestException:
                all_results.append({"link": link, "text": None})
    return all_results


def respond(message, history):
    if history is None:
        history = []

    messages = []

    functions_metadata = [
        {"type": "function", "function": {"name": "web_search", "description": "Search query on google", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "web search query"}}, "required": ["query"]}}},
    ]

    system_messages = f'[SYSTEM]You are a helpful assistant, and respond without emoji, You have access to the following functions: \n {str(functions_metadata)}\n\nTo use these functions respond with:\n<functioncall> {{ "name": "function_name", "arguments": {{ "arg_1": "value_1", "arg_1": "value_1", ... }} }}  </functioncall>  [USER] {message}'
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=512,
        stream=True,
        temperature=0.1,
        top_p=0.95,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response


async def tts(response):
    tts = edge_tts.Communicate(response, voice="en-US-JennyNeural")
    await tts.save("output.mp3")
    return "output.mp3"


def generate_tts(response):
    return asyncio.run(tts(response))


def process(message, history):
    response = list(respond(message, history))[-1]
    tts_file = generate_tts(response)
    history.append((message, response))
    return history, history, tts_file


with gr.Blocks() as demo:
    gr.Markdown("<h1><center>gemma-2b-it</center></h1>")

    chatbot = gr.Chatbot()
    audio = gr.Audio(label='Your audio output')

    input_text = gr.Textbox(placeholder="Type a message...")
    submit_button = gr.Button("Submit")

    submit_button.click(
        process,
        inputs=[input_text, chatbot],
        outputs=[chatbot, chatbot, audio]
    )

if __name__ == "__main__":
    demo.launch()