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
import yagmail

client = openai.OpenAI()

assistant_id = "asst_NHnYFIdpvioacAJqWYMchJHI"
vector_id = "vs_sqT4VRRTwkH7JPr3AT8CpoXV"

def send_email(thread):
    mail_from = os_getenv("MAIL_FROM")
    mail=yagmail.SMTP(
        mail_from,
        os_getenv("MAIL_PASS"),
        host=os_getenv("MAIL_HOST"),
        port=int(os_getenv("MAIL_PORT"))
    )
    mail.send(mail_from,f"New thread {thread.id}")

    
def chat(user_message, history, state):
    if (state is None) or (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
            send_email(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,
        ) as stream:
            total = ''
            for delta in stream.text_deltas:
                total += delta
                yield total

def chat_nostream(user_message, history, state):
    if (state is None) or (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

        # 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", fill_height=True) 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=[
            ["Qué diferencias hay entre el cuerpo y el Alma?"],
            ["Cuáles son las distintas funciones de los órganos del cuerpo y su relación con 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)