import os, threading
import gradio as gr
from crew import run_crew
from utils import get_questions

QUESTION_FILE_PATH = "data/gaia_validation.jsonl"
QUESTION_LEVEL     = 1

def invoke(level, question, file_name, ground_truth, openai_api_key, gemini_api_key):
    if not question:
        raise gr.Error("Question is required.")

    if not openai_api_key:
        raise gr.Error("OpenAI API Key is required.")
    
    if not gemini_api_key:
        raise gr.Error("Gemini API Key is required.")

    if file_name:
        file_name = f"data/{file_name}"
    
    lock = threading.Lock()
    
    with lock:
        answer = ""

        try:
            os.environ["OPENAI_API_KEY"] = openai_api_key
            os.environ["GEMINI_API_KEY"] = gemini_api_key
            
            answer = run_crew(question, file_name)
        except Exception as e:
            raise gr.Error(e)
        finally:
            del os.environ["OPENAI_API_KEY"]
            del os.environ["GEMINI_API_KEY"]
        
        return answer

gr.close_all()

with gr.Blocks() as gaia:
    gr.Markdown("# General AI Assistant - GAIA 🤖🤝🤖")
    gr.Markdown(os.environ["DESCRIPTION"])
    gr.Radio([1, 2, 3], label="Level", value=int(os.environ["INPUT_LEVEL"]))
    gr.Markdown(label="Question *", value=os.environ["INPUT_QUESTION"])
    gr.Textbox(label="File Name")
    gr.Textbox(label="Ground Truth", value=os.environ["INPUT_GROUND_TRUTH"])
    with gr.Row():
        gr.Textbox(label="OpenAI API Key *", type="password")
        gr.Textbox(label="Gemini API Key *", type="password")
    gr.Textbox(label="Answer", lines=1, interactive=False, value=os.environ["OUTPUT"])
    submit_btn = gr.Button("Submit")   
    submit_btn.click(
        fn=invoke,
        inputs=[level, question, filename, ground_truth, openai_key, gemini_key],
        outputs=answer
    )

gaia.launch()