import time
import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv
from pathlib import Path
import json
load_dotenv()

# client = OpenAI()
client = OpenAI()

assistant_id = 'asst_4DJlWbqXooIm6hzJYBW5TSpX'


def wait_on_run(run, thread):
    while run.status == "queued" or run.status == "in_progress":
        run = client.beta.threads.runs.retrieve(
            thread_id=thread,
            run_id=run.id,
        )
        time.sleep(0.5)
    return run

# thread = client.beta.threads.create()
# print("first print thread ID", thread.id)


def send_message(message, thread_id):
    message = client.beta.threads.messages.create(thread_id, role='user', content=message)
    run = client.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id=assistant_id,
    )
    run = wait_on_run(run, thread_id)
    messages = client.beta.threads.messages.list(thread_id=thread_id)
    if run.status == "requires_action":
        tool_call = run.required_action.submit_tool_outputs.tool_calls[0]
        # name = tool_call.function.name
        # arguments = json.loads(tool_call.function.arguments)
        run = client.beta.threads.runs.submit_tool_outputs(
                thread_id=thread_id,
                run_id=run.id,
                tool_outputs=[
                    {"tool_call_id": tool_call.id, "output": json.dumps("true")}
                ]
            )
        
        try:
            print("second print thread_id", thread_id)
            run = wait_on_run(run, thread_id)
        except:
            # print("Third print thread", thread)
            # run = wait_on_run(run, thread)
            print("Third print thread", thread_id)
    messages = client.beta.threads.messages.list(thread_id=thread_id)
    last_message = messages.data[0].content[0].text.value
    return last_message

def get_thread_id():
    thread = client.beta.threads.create()
    print("first print thread ID", thread.id)
    return thread.id

def file_update_for_threads(filepath, user_message):
    filename = filepath.name
    file = client.files.create(
        file=open(filename, "rb"),
        purpose="assistants",
    )
    thread = client.beta.threads.create(
        messages=[
            {
            "role": "user",
            "content": user_message,
            "attachments": [
                { "file_id": file.id, "tools": [{"type": "file_search"}] }
            ],
            }
        ]
    )
    run = client.beta.threads.runs.create_and_poll(
        thread_id=thread.id, assistant_id=assistant_id
    )
    messages = list(client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id))
    message_content = messages[0].content[0].text.value
    return message_content, thread.id

with gr.Blocks() as demo:
    with gr.Row():
        file_upload = gr.File(label="Upload your File")
        user_message_input = gr.Textbox(label="Enter your message")
    with gr.Row():
        output_message = gr.Textbox(label="Response", interactive=False)
        output_thread_id = gr.Textbox(label="Thread ID", interactive=False)
    upload_button = gr.Button("Upload File and Send Message")
    
    upload_button.click(
        file_update_for_threads, 
        inputs=[file_upload, user_message_input], 
        outputs=[output_message, output_thread_id]
    )
    # get_thread = gr.Textbox(label="Get Thread ID")
    # get_thread_d = gr.Button("Get Thread ID")
    # get_thread_d.click(fn=get_thread_id, outputs=get_thread, api_name="get_thread_id")
    message_to_assistant = gr.Textbox(label="Send Message to the the current thread")
    thread_id = gr.Textbox(label="Thread ID")
    output = gr.Textbox(label="Output response from the Assistant")
    greet_btn = gr.Button("Send Message")
    greet_btn.click(fn=send_message, inputs=[message_to_assistant, thread_id], outputs=output, api_name="send_message")

demo.launch()