from openai import OpenAI, AssistantEventHandler
from dotenv import load_dotenv
import os
import gradio as gr

# Load environment variables from .env file
load_dotenv()

# Load env variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ASSISTANT_ID = os.getenv("ASSISTANT_ID")

client = OpenAI(api_key=OPENAI_API_KEY)

def create_thread():
    return client.beta.threads.create()

def predict(user_message, history, thread):
    # Append the new user message to the history
    # history.append({"role": "user", "content": user_message})

    # Send the user message to the OpenAI API
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=user_message
    )

    response = []
    # Stream the assistant's response
    with client.beta.threads.runs.stream(
        thread_id=thread.id,
        assistant_id=ASSISTANT_ID
    ) as stream:
        for event in stream:
            if event.event == "thread.message.delta" and event.data.delta.content:
                assistant_message = event.data.delta.content[0].text
                response.append(assistant_message.value)
                yield [{"role": "assistant", "content": ''.join(response)}], thread
    
    # # Append the assistant's response to the history
    history.append({"role": "assistant", "content": ''.join(response)})
    yield history, thread

# Launch the Gradio chat interface
with gr.Blocks(title="Hair Library Shopping Assistant Demo") as demo:
    chatbot = gr.Chatbot(type='messages', label="Hair Library Shopping Assistant Demo")
    msg = gr.Textbox(placeholder="Type your message here...")
    user_content = gr.State("")
    thread = gr.State(create_thread)
    history = gr.State([])
    submit = gr.Button("Submit")

    def user_input(user_message, history):
        user_content = user_message
        return "", history + [{"role": "user", "content": user_message}], user_content

    msg.submit(user_input, [msg, chatbot], [msg, chatbot, user_content], queue=False).then(
        predict, [user_content, chatbot, thread], [chatbot, thread]
    )
    submit.click(user_input, [msg, chatbot], [msg, chatbot, user_content], queue=False).then(
        predict, [user_content, chatbot, thread], [chatbot, thread]
    )

    # msg.submit(predict, [msg, chatbot, thread], [chatbot, thread], queue=False)
    # clear.click(lambda: None, None, chatbot, queue=False)

demo.launch(share=False)