import os
import gradio as gr
from agents.crew import run_crew
from utils.utils import get_questions

def ask(question, openai_api_key, gemini_api_key, anthropic_api_key, file_name = ""):
    """
    Ask General AI Assistant a question to answer.

    Args:
        question (str): The question to answer
        openai_api_key (str): OpenAI API key
        gemini_api_key (str): Gemini API key
        anthropic_api_key (str): Anthropic API key
        file_name (str): Optional file name

    Returns:
        str: The answer to the question
    """
    if not question:
        gr.Warning("Question is required.")
        return ""

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

    if not anthropic_api_key:
        gr.Warning("Anthropic API Key is required.")
        return ""
        
    if file_name:
        file_name = f"data/{file_name}"
    
    try:
        os.environ["OPENAI_API_KEY"] = openai_api_key
        os.environ["GEMINI_API_KEY"] = gemini_api_key
        os.environ["MODEL_API_KEY"] = anthropic_api_key
        
        return run_crew(question, file_name)
    except Exception as e:
        gr.Warning(str(e))
        return ""
    finally:
        for key in ["OPENAI_API_KEY", "GEMINI_API_KEY", "MODEL_API_KEY"]:
            os.environ.pop(key, None)

with gr.Blocks() as gaia:
    gr.Markdown("## GAIA - General AI Assistant")

    gr.Markdown(os.environ.get("DESCRIPTION"))

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                question = gr.Textbox(
                    label="Question *",
                    placeholder="In \"MCP's 1st Birthday Hackathon - Hosted by Anthropic and Gradio\", what percentage of participants submitted a solution so far?",
                    value="Black to move. Without moving the black queens, which sequence is mate in 2 for black, regardless of what white does? Use standard chess notation, leaving out the white move.",
                    interactive=True
                )
            with gr.Row():
                level = gr.Radio(
                    choices=[1, 2, 3],
                    label="GAIA Benchmark Level",
                    interactive=True,
                    scale=1
                )
                ground_truth = gr.Textbox(
                    label="Ground Truth",
                    value="Rxf3, Rf1#",
                    interactive=True,
                    scale=1
                )
                file_name = gr.Textbox(
                    label="File Name",
                    value="hle-6687ffb1091058ff19128813.jpg",
                    interactive=True
                )
            with gr.Row():
                openai_api_key = gr.Textbox(
                    label="OpenAI API Key *",
                    type="password",
                    placeholder="sk‑...",
                    interactive=True
                )
                gemini_api_key = gr.Textbox(
                    label="Gemini API Key *",
                    type="password",
                    interactive=True
                )
                anthropic_api_key = gr.Textbox(
                    label="Anthropic API Key *",
                    type="password",
                    placeholder="sk-ant-...",
                    interactive=True
                )
            with gr.Row():
                clear_btn = gr.ClearButton(
                    components=[question, level, ground_truth, file_name]
                )
                submit_btn = gr.Button("Submit", variant="primary")
        with gr.Column(scale=1):
            answer = gr.Textbox(
                label="Answer",
                interactive=False
            )

    submit_btn.click(
        fn=ask,
        inputs=[question, openai_api_key, gemini_api_key, anthropic_api_key, file_name],
        outputs=answer
    )

    QUESTION_FILE_PATH = "data/gaia_validation.jsonl"
    
    gr.Examples(
        label="GAIA Benchmark Level 1 Problems",
        examples=get_questions(QUESTION_FILE_PATH, 1),
        inputs=[question, level, ground_truth, file_name, openai_api_key, gemini_api_key, anthropic_api_key],
        outputs=answer,
        cache_examples=False
    )

    gr.Examples(
        label="GAIA Benchmark Level 2 Problems",
        examples=get_questions(QUESTION_FILE_PATH, 2),
        inputs=[question, level, ground_truth, file_name, openai_api_key, gemini_api_key, anthropic_api_key],
        outputs=answer,
        cache_examples=False
    )

    gr.Examples(
        label="GAIA Benchmark Level 3 Problems",
        examples=get_questions(QUESTION_FILE_PATH, 3),
        inputs=[question, level, ground_truth, file_name, openai_api_key, gemini_api_key, anthropic_api_key],
        outputs=answer,
        cache_examples=False
    )

gaia.launch(mcp_server=True, share=False, ssr_mode=False)