"""Unit 4: fetch the 20 GAIA questions, answer them with the agent, submit for the certificate.

Local:  py -3.12 app.py             -> answer all questions (cached in answers.json)
        py -3.12 app.py --only 3 9  -> (re)answer only those question numbers
        py -3.12 app.py --submit    -> submit the cached answers and print the score
Space:  same code with a Gradio login + button (runs when SPACE_ID is set).
"""
import json
import os
import sys
from pathlib import Path

import requests

API = "https://agents-course-unit4-scoring.hf.space"
USERNAME = "brunopaschoal"
AGENT_CODE = "https://huggingface.co/spaces/brunopaschoal/Final_Assignment/tree/main"
HERE = Path(__file__).parent
CACHE = HERE / "answers.json"
FILES = HERE / "files"
TEXT_FILES = (".py", ".txt", ".csv", ".md", ".json")


def get_questions() -> list[dict]:
    r = requests.get(f"{API}/questions", timeout=30)
    r.raise_for_status()
    return r.json()


def build_prompt(q: dict) -> str:
    prompt = q["question"]
    if not q.get("file_name"):
        return prompt
    FILES.mkdir(exist_ok=True)
    path = FILES / q["file_name"]
    if not path.exists():
        r = requests.get(f"{API}/files/{q['task_id']}", timeout=60)
        if r.ok:
            path.write_bytes(r.content)
        else:  # course server often 404s; same file is in the (gated) GAIA validation set
            from huggingface_hub import hf_hub_download
            hf_hub_download("gaia-benchmark/GAIA", f"2023/validation/{q['file_name']}", repo_type="dataset", local_dir=HERE / "gaia")
            path.write_bytes((HERE / "gaia/2023/validation" / q["file_name"]).read_bytes())
    prompt += f"\n\nAttached file: {path.resolve().as_posix()}"
    if path.suffix in TEXT_FILES:  # the agent's sandbox can't open() files, so inline text files
        prompt += f"\nFile contents:\n```\n{path.read_text(encoding='utf-8', errors='replace')}\n```"
    return prompt


def answer_all(only: set[int] | None = None) -> dict:
    from agent import build_agent, clean_answer

    cache = json.loads(CACHE.read_text(encoding="utf-8")) if CACHE.exists() else {}
    for i, q in enumerate(get_questions(), 1):
        if (only and i not in only) or (not only and q["task_id"] in cache):
            continue
        print(f"\n===== Question {i}/20 =====\n{q['question'][:200]}")
        try:
            cache[q["task_id"]] = clean_answer(build_agent().run(build_prompt(q)))
        except Exception as e:  # one failure (e.g. quota) must not lose the other answers
            print(f"!! Question {i} failed: {e}")
            continue
        CACHE.write_text(json.dumps(cache, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f">> Answer {i}: {cache[q['task_id']]}")
    return cache


def submit(cache: dict, username: str = USERNAME) -> str:
    answers = [{"task_id": t, "submitted_answer": a} for t, a in cache.items()]
    r = requests.post(f"{API}/submit", json={"username": username, "agent_code": AGENT_CODE, "answers": answers}, timeout=120)
    r.raise_for_status()
    d = r.json()
    return f"Score: {d.get('score')}% ({d.get('correct_count')}/{d.get('total_attempted')} correct). {d.get('message', '')}"


def launch_space():
    import gradio as gr

    def run(profile: gr.OAuthProfile | None):
        if not profile:
            return "Please log in with Hugging Face first."
        return submit(answer_all(), profile.username)

    with gr.Blocks() as demo:
        gr.Markdown("# GAIA agent (HF Agents course, Unit 4)\nsmolagents CodeAgent + Gemini. Log in, then run.")
        gr.LoginButton()
        out = gr.Textbox(label="Result")
        gr.Button("Run evaluation & submit").click(run, outputs=out)
    demo.launch()


if __name__ == "__main__":
    if os.environ.get("SPACE_ID"):
        launch_space()
    elif "--submit" in sys.argv:
        print(submit(json.loads(CACHE.read_text(encoding="utf-8"))))
    else:
        only = {int(a) for a in sys.argv[sys.argv.index("--only") + 1:]} if "--only" in sys.argv else None
        answer_all(only)