from openai.types.beta.threads import Message, MessageDelta
from os import getenv as os_getenv, path as os_path
from time import sleep
from json import loads as json_loads
import gradio as gr
import openai

OPENAI_API_KEY = os_getenv('OPENAI_API_KEY')

client = openai.OpenAI(api_key=OPENAI_API_KEY)

assistant_id = "asst_NHnYFIdpvioacAJqWYMchJHI"
vector_id = "vs_sqT4VRRTwkH7JPr3AT8CpoXV"

class EventHandler(openai.AssistantEventHandler):
    # def on_event(self, event):
    #     print(f"event: {event.event}\n{event.data}\n\n")
        
    def on_tool_call_created(self, tool_call):
        print(f"\nassistant > {tool_call.type}\n", flush=True)
    
    def on_tool_call_delta(self, delta, snapshot):
        if delta.type == 'code_interpreter':
            if delta.code_interpreter.input:
                print(delta.code_interpreter.input, end="", flush=True)
            if delta.code_interpreter.outputs:
                print(f"\n\noutput >", flush=True)
                for output in delta.code_interpreter.outputs:
                    if output.type == "logs":
                        print(f"\n{output.logs}", flush=True)

def chat(user_message, history, state):
    if not state['user']:
        gr.Warning("You need to authenticate first")
        yield
    else:
        thread = state['thread']
        if thread is None:
            thread = client.beta.threads.create(
                tool_resources={
                    "file_search": {
                        "vector_store_ids": [vector_id]
                    }
                }
            )
            state['thread'] = thread

        client.beta.threads.messages.create(
            thread_id=thread.id,
            role="user",
            content=user_message,
        )

        with client.beta.threads.runs.stream(
            thread_id=thread.id,
            assistant_id=assistant_id,
            event_handler=EventHandler(),
        ) as stream:
            total = ''
            for delta in stream.text_deltas:
                total += delta
                yield total

def chat_nostream(user_message, history, state):
    if state['user'] is None:
        return

    thread = state['thread']
    if thread is None:
        thread = client.beta.threads.create(
            tool_resources={
                "file_search": {
                    "vector_store_ids": [vector_id]
                }
            }
        )
        state['thread'] = thread

    # Add the user's message to the thread
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=user_message,
    )

    # Run the Assistant
    run = client.beta.threads.runs.create(thread_id=thread.id,
                                        assistant_id=assistant_id)

    while True:
        run_status = client.beta.threads.runs.retrieve(thread_id=thread.id,
                                                        run_id=run.id)
        print(f"Run status: {run_status.status}")
        if run_status.status == 'completed':
            break
        sleep(5)

    messages = client.beta.threads.messages.list(thread_id=thread.id)
    response = messages.data[0].content[0].text.value

    yield response

def new_state():
    return gr.State({
        "user": None,
        "thread": None,
    })

def auth(token, state):
    tokens=os_getenv("APP_TOKENS", None)
    if tokens is None:
        state["user"] = "anonymous"
    else:
        tokens=json_loads(tokens)
        state["user"] = tokens.get(token, None)
    return "", state

AUTH_JS = """function auth_js(token, state) {
    if (!!document.location.hash) {
        token = document.location.hash
        document.location.hash=""
    }        
    return [token, state]
}
"""
with gr.Blocks(title="Dr. Luis Chiozza - Medicina y Psicoanalisis") as demo:
    state = new_state()
    gr.ChatInterface(
        chat,
        title="Dr. Luis Chiozza - Medicina y Psicoanalisis",
        description="Habla con la colección de Medicina y Psicoanalisis del Dr. Luis Chiozza",
        additional_inputs=[state],
        examples=[["Que diferencias hay entre el cuerpo y el Alma?"]],
    )

    token = gr.Textbox(visible=False)
    demo.load(auth,
        [token,state],
        [token,state],
        js=AUTH_JS)

demo.launch(
    height=700,
    allowed_paths=["."])
    # auth_dependency=authenticate) #auth=authenticate)