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(question, level, 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()

"""
gaia = gr.Interface(fn=invoke, 
                    inputs=[gr.Radio([1, 2, 3], label="Level", value=int(os.environ["INPUT_LEVEL"])),
                            gr.Textbox(label="Question *", value=os.environ["INPUT_QUESTION"]),
                            gr.Textbox(label="File Name"),
                            gr.Textbox(label="Ground Truth", value=os.environ["INPUT_GROUND_TRUTH"]),
                            gr.Textbox(label="OpenAI API Key *", type="password"),
                            gr.Textbox(label="Gemini API Key *", type="password")],
                    outputs=[gr.Textbox(label="Answer", lines=1, interactive=False, value=os.environ["OUTPUT"])],
                    title="General AI Assistant - GAIA 🤖🤝🤖",
                    description=os.environ["DESCRIPTION"],
                    examples=get_questions(QUESTION_FILE_PATH, QUESTION_LEVEL),
                    cache_examples=False
    )

gaia.launch()
"""

examples = get_questions(QUESTION_FILE_PATH, QUESTION_LEVEL)

with gr.Blocks(title="GAIA") as gaia:
    gr.Markdown("## General AI Assistant - GAIA 🤖🤝🤖")
    gr.Markdown(os.environ.get("DESCRIPTION", ""))

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                question = gr.Textbox(
                    label="Question *",
                    value=os.environ.get("INPUT_QUESTION", "")
                )
            with gr.Row():
                with gr.Column(scale=1):
                    level = gr.Radio(
                        choices=[1, 2, 3],
                        label="Level",
                        value=int(os.environ.get("INPUT_LEVEL", 1))
                    )
                with gr.Column(scale=1):
                    ground_truth = gr.Textbox(
                        label="Ground Truth",
                        value=os.environ.get("INPUT_GROUND_TRUTH", "")
                    )
                with gr.Column(scale=2):
                    file_name = gr.Textbox(label="File Name")
            with gr.Row():
                openai_api_key = gr.Textbox(
                    label="OpenAI API Key *",
                    type="password"
                )
                gemini_api_key = gr.Textbox(
                    label="Gemini API Key *",
                    type="password"
                )
            with gr.Row():
                clear_btn = gr.ClearButton(
                    components=[question, level, file_name, ground_truth, openai_api_key, gemini_api_key],
                    value="Clear"
                )
                submit_btn = gr.Button("Submit", variant="primary")
        with gr.Column(scale=1):
            answer = gr.Textbox(
                label="Answer",
                lines=1,
                interactive=False,
                value=os.environ.get("OUTPUT", "")
            )

    submit_btn.click(
        fn=invoke,
        inputs=[question, level, file_name, ground_truth, openai_api_key, gemini_api_key],
        outputs=answer
    )

    gr.Examples(
        examples=examples,
        inputs=[question, level, file_name, ground_truth, openai_api_key, gemini_api_key],
        outputs=answer,
        fn=invoke,
        cache_examples=False
    )

gaia.launch()