import gradio as gr
import json
import os
from transformers import pipeline

# Load model (lightweight & allowed)
qa_pipeline = pipeline(
    "text2text-generation",
    model="google/flan-t5-base",
    max_new_tokens=64
)

DATA_PATH = "/data/gaia_validation_questions.json"


def solve_question(question: str) -> str:
    """
    Very simple baseline solver.
    GAIA Unit 4 rewards formatting + correctness, not fancy agents.
    """
    try:
        result = qa_pipeline(question)
        return result[0]["generated_text"].strip()
    except Exception:
        return "unknown"


def run_evaluation():
    """
    Runs GAIA evaluation and returns answers in correct format
    """
    with open(DATA_PATH, "r") as f:
        data = json.load(f)

    answers = {}

    for item in data:
        qid = item["id"]
        question = item["question"]
        answers[qid] = solve_question(question)

    return answers


def submit():
    """
    This function is REQUIRED by GAIA.
    It must return a dict of {question_id: answer}
    """
    return run_evaluation()


with gr.Blocks() as demo:
    gr.Markdown("# GAIA Unit 4 – Basic Agent Runner")
    run_btn = gr.Button("Run Evaluation & Submit")
    output = gr.JSON(label="Submission Result")

    run_btn.click(fn=submit, outputs=output)

demo.launch(server_name="0.0.0.0", server_port=7860)