import gradio as gr
import json, openai, os, time

from openai import OpenAI

_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

_assistant, _thread = None, None

def show_json(str, obj):
    print(f"=> {str}\n{json.loads(obj.model_dump_json())}")

def init_assistant(client):
    assistant = client.beta.assistants.create(
        name="Math Tutor",
        instructions="You are a personal math tutor. Answer questions briefly, in a sentence or less.",
        model="gpt-4-1106-preview",
    )
    show_json("assistant", assistant)
    return assistant
    
def init_thread(client):
    thread = client.beta.threads.create()
    show_json("thread", thread)
    return thread
    
def wait_on_run(client, thread, run):
    while run.status == "queued" or run.status == "in_progress":
        run = client.beta.threads.runs.retrieve(
            thread_id=thread.id,
            run_id=run.id,
        )
        time.sleep(0.25)
    return run

def extract_content_values(data):
    content_values = []
    for item in data.data:
        for content in item.content:
            if content.type == 'text':
                content_values.append(content.text.value)
    return content_values

def chat(message, history, openai_api_key):
    global _client, _assistant, _thread     
       
    if _client == None:
        _client = init_client()

    if _assistant == None:
        _assitant = init_assistant(_client)

    if _thread == None:
        _thread = init_thread(_client)
        
    message = _client.beta.threads.messages.create(
        role="user",
        thread_id=_thread.id,
        content=message,
    )
    show_json("message", message)
    
    run = _client.beta.threads.runs.create(
        assistant_id=_assistant.id,
        thread_id=_thread.id,
    )
    show_json("run", run)
    
    run = wait_on_run(_client, _thread, run)
    show_json("run", run)
    
    messages = _client.beta.threads.messages.list(
        thread_id=_thread.id
    )
    show_json("messages", messages)

    return extract_content_values(messages)[0]
        
gr.ChatInterface(
    chat,
    chatbot=gr.Chatbot(height=300),
    textbox=gr.Textbox(placeholder="Question", container=False, scale=7),
    title="Multi-Assistant Demo",
    description="Ask AAA Assistant, BBB Assistant, and CCC Assistant any question",
    retry_btn=None,
    undo_btn=None,
    clear_btn="Clear",
    examples=[["I need to solve the equation '2x + 10 = 20'. Can you help me?", "sk-<BringYourOwn>"]],
    cache_examples=False,
    additional_inputs=[
        gr.Textbox("sk-", label="OpenAI API Key", type = "password"),
    ],
).launch()