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 "Submit your question as input !"
    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 == "Examples":
        filename = example[-6:-2] + ".md"
        file = open("examples/" + filename, "r")
        output = file.read()
        yield output
    else:
        # check the password
        if password == "":
            yield "To search you need to enter an RTL password !"
        elif password != mypassword:
            yield "Please enter the correct RTL password !"
        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 Multilingual AI News Reader : What happened in the country 🇱🇺 or in the world 🌎 ?</h2>")   
    with gr.Row():
        myDescription = gr.HTML("""
            <h3 align='center'>What topic interests you ?</h3>
            <p align='center'>🐶 🏃🏻‍♂️ 🌗 🍇 🌈 🍽️ 🏆 🚘 ✈️ 🩺 </p>
            <p align='center' bgcolor="Moccasin">Submit your question in english or in another language !</p>
            """            
        )
    with gr.Row():
        mode = gr.Radio(choices=["Search", "Examples"], label = "You can run the examples without password !", value = "Examples")
        pw = gr.Textbox(lines=1, label="Enter the RTL password : ")
    with gr.Row():
            question = gr.Textbox(lines=3, label="Please submit your question ?")
    with gr.Row():    
            examples = gr.Radio(["What is reported about animals in the news ?", "Quelles sont les rubriques régulières dans les informations de RTL ?", "Wat ware Wierder vum Joer an der Vergaangenheet ?", "Welches Thema kommt sehr häufig vor in den RTL Nachrichten ?"], value="Wat ware Wierder vum Joer an der Vergaangenheet ?", label="Examples")
    with gr.Row():
        clear = gr.Button("Clear")
        submit = gr.Button("Submit")
    with gr.Row():
        mySubtitle = gr.HTML("<p align='center' bgcolor='Khaki'>English RTL News :</p>")
    with gr.Row():
        myOutput = gr.Markdown(label="Answer from the OpenAI File-Search Assistent :")
    # with gr.Row():
    #    myInfo = gr.HTML("<p>Please request your RTL password per mail from <b>marco.barnig@pt.lu<b></p>")

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