from openai import AsyncAssistantEventHandler
from openai import AsyncOpenAI
import gradio as gr
import asyncio
import os

# set the keys
client = AsyncOpenAI(
  api_key=os.getenv("OPENAI_API_KEY")
)

assistantID = os.getenv("OPENAI_ASSISTANT_ID")
mypassword = os.getenv("RTL_PASSWORD")

class EventHandler(AsyncAssistantEventHandler):
    def __init__(self) -> None:
        super().__init__()
        self.response_text = ""

    async def on_text_created(self, text) -> None:
        self.response_text += str(text)

    async def on_text_delta(self, delta, snapshot):
        self.response_text += str(delta.value)

    async def on_text_done(self, text):
        pass

    async def on_tool_call_created(self, tool_call):
        self.response_text += f"\n[Tool Call]: {str(tool_call.type)}\n"

    async def on_tool_call_delta(self, delta, snapshot):
        if snapshot.id != getattr(self, "current_tool_call", None):
            self.current_tool_call = snapshot.id
            self.response_text += f"\n[Tool Call Delta]: {str(delta.type)}\n"

        if delta.type == 'code_interpreter':
            if delta.code_interpreter.input:
                self.response_text += str(delta.code_interpreter.input)
            if delta.code_interpreter.outputs:
                self.response_text += "\n\n[Output]:\n"
                for output in delta.code_interpreter.outputs:
                    if output.type == "logs":
                        self.response_text += f"\n{str(output.logs)}"

    async def on_tool_call_done(self, text):
        pass

# Initialize session variables
session_data = {"assistant_id": assistantID, "thread_id": None}

async def initialize_thread():
    # Create a Thread
    thread = await client.beta.threads.create()
    # Store thread ID in session_data for later use
    session_data["thread_id"] = thread.id

async def generate_response(user_input):
    if user_input == "":
        yield "Veuillez soumettre une question !"
    else:
        assistant_id = session_data["assistant_id"]
        thread_id = session_data["thread_id"]

        # Add a Message to the Thread
        oai_message = await client.beta.threads.messages.create(
            thread_id=thread_id,
            role="user",
            content=user_input
    )

    # Create and Stream a Run
    event_handler = EventHandler()

    async with client.beta.threads.runs.stream(
        thread_id=thread_id,
        assistant_id=assistant_id,
        instructions="Please assist the user with their query.",
        event_handler=event_handler,
    ) as stream:
        # Yield incremental updates
        async for _ in stream:
            await asyncio.sleep(0.1)  # Small delay to mimic streaming
            yield event_handler.response_text

# Gradio interface function (generator)
async def gradio_chat_interface(mode, password, user_input, example):
    if mode == "Exemples":
        filename = example[-6:-2] + ".md"
        file = open("examples/" + filename, "r")
        output = file.read()
        yield output
    else:
        # check the password
        if password == "":
            yield "Pour faire des recherches vous devez entrer un mot de passe !"
        elif password != mypassword:
            yield "Veuillez entrer le mot de passe correct !"
        else:

            # Create a new event loop if none exists (or if we are in a new thread)
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)

            # Initialize the thread if not already done
            if session_data["thread_id"] is None:
                await initialize_thread()

            # Generate and yield responses
            async for response in generate_response(user_input):
                yield response

with gr.Blocks() as demo:
    with gr.Row():
        myTitle = gr.HTML("<h2 align=center>RTL French AI News Reader : Quels événements ont eu lieu dans le Grand-Duché 🇱🇺 ou dans le monde 🌎 ?</h2>")   
    with gr.Row():
        myDescription = gr.HTML("""
            <h3 align='center'>Quel sujet vous intéresse ?</h3>
            <p align='center'>🐶 🏃🏻‍♂️ 🌗 🍇 🌈 🍽️ 🏆 🚘 ✈️ 🩺 </p>
            <p align='center' bgcolor="Moccasin">Veuillez soumettre votre question en français ou dans une autre langue !</p>
            """            
        )
    with gr.Row():
        mode = gr.Radio(choices=["Recherche", "Exemples"], label = "Vous pouvez lire les exemples sans mots de passe !", value = "Exemples")
        pw = gr.Textbox(lines=1, label="Veuillez entrer le mot de passe :")
    with gr.Row():
            question = gr.Textbox(lines=3, label="Veuillez soumettre votre question :")
    with gr.Row():    
            examples = gr.Radio(["Prière d'établir la liste des meilleurs sportifs luxembourgeois votés par le public entre 2014 et 2023 !", "Quelle évolution a eu lieu dans le domaine de la protection de la nature pendant les dix dernières années ?"], value="Prière d'établir la liste des meilleurs sportifs luxembourgeois votés par le public entre 2014 et 2023 !", label="Exemples")
    with gr.Row():
        clear = gr.Button("Clear")
        submit = gr.Button("Submit")
    with gr.Row():
        mySubtitle = gr.HTML("<p align='center' bgcolor='Khaki'>French RTL News :</p>")
    with gr.Row():
        myOutput = gr.Markdown(label="Réponses de l'assistant OpenAI File-Search :")

    submit.click(fn = gradio_chat_interface, inputs=[mode, pw, question, examples], outputs = myOutput)
demo.launch()