import requests
import gradio as gr
import os
import json
from agent import SimpleAgent

# --- CONFIG: set this to the official course API base if different ---
API_BASE = os.environ.get("GAIA_API_BASE", "https://huggingface-course-api.hf.space") 
# (replace above if course gave another URL)

# instantiate agent
agent = SimpleAgent()  # loads model (may take a few seconds)

def get_questions():
    """
    GET /questions -> returns list of tasks (expected fields: task_id, question, ...)
    """
    url = f"{API_BASE}/questions"
    r = requests.get(url, timeout=15)
    r.raise_for_status()
    return r.json()

def get_random_question():
    url = f"{API_BASE}/random-question"
    r = requests.get(url, timeout=10)
    r.raise_for_status()
    return r.json()

def run_agent_on_questions(questions):
    answers = []
    for q in questions:
        task_id = q.get("task_id")
        question_text = q.get("question") or q.get("task") or q.get("input", "")
        # Ensure we send only the question string
        ans = agent.answer(question_text)
        answers.append({"task_id": task_id, "submitted_answer": ans})
    return answers

def run_and_submit(username, code_link):
    """
    Performs: fetch questions, run agent, submit payload to /submit
    Returns API response.
    """
    if not username or not code_link:
        return {"error": "Please provide both Hugging Face username and public code link."}
    # 1. fetch questions
    questions = get_questions()
    # 2. run agent
    answers = run_agent_on_questions(questions)
    # 3. build payload
    payload = {
        "username": username,
        "agent_code": code_link,
        "answers": answers
    }
    # 4. submit
    url = f"{API_BASE}/submit"
    r = requests.post(url, json=payload, timeout=30)
    try:
        return r.json()
    except Exception:
        return {"status_code": r.status_code, "text": r.text}

def preview_and_run_one():
    q = get_random_question()
    ans = agent.answer(q.get("question", ""))
    return f"Task id: {q.get('task_id')}\nQuestion:\n{q.get('question')}\n\nAgent answer:\n{ans}"

# GRADIO UI
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Unit-4 Agent Runner (Course Template)\nRun agent on the filtered GAIA questions and submit to leaderboard.")
    with gr.Row():
        hf_user = gr.Textbox(label="Hugging Face username", placeholder="your-hf-username")
        code_link = gr.Textbox(label="Public link to your Space code (tree/main)", placeholder="https://huggingface.co/spaces/<username>/<space>/tree/main")
    run_btn = gr.Button("Run agent & Submit")
    output = gr.JSON()

    run_btn.click(fn=run_and_submit, inputs=[hf_user, code_link], outputs=output)

    gr.Markdown("## Quick test")
    test_btn = gr.Button("Run agent on one random question (preview)")
    test_out = gr.Textbox()
    test_btn.click(fn=preview_and_run_one, inputs=None, outputs=test_out)

if __name__ == "__main__":
    demo.launch()