# app.py — GAIA Benchmark Agent (FastAPI)
import os
import json
import tempfile
import logging
from pathlib import Path

import requests
from fastapi import FastAPI, Request, Form, File, UploadFile, BackgroundTasks
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.templating import Jinja2Templates

from agent import build_agent, gaia_benchmark_iter

HF_TOKEN    = os.environ.get("HF_TOKEN", "")
COURSE_API  = "https://agents-course-unit4-scoring.hf.space"
SPACE_URL   = "https://huggingface.co/spaces/gilra/wreck/tree/main"
HF_USERNAME = "gilra"

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

app = FastAPI(title="GAIA Benchmark Agent")
templates = Jinja2Templates(directory="templates")


# ── Single-question endpoint (existing) ─────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse(request=request, name="index.html")


@app.post("/ask")
async def ask(
    question: str = Form(...),
    file: UploadFile = File(None),
):
    if not HF_TOKEN:
        return JSONResponse({"error": "HF_TOKEN secret not set on this Space."}, status_code=500)

    tmp_path = None
    if file and file.filename:
        suffix = Path(file.filename).suffix or ".bin"
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(await file.read())
            tmp_path = tmp.name

    full_q = question.strip()
    if tmp_path:
        full_q += f"\n\nAttached file path: {tmp_path}"

    agent = build_agent(HF_TOKEN)
    try:
        answer = str(agent.run(full_q))
    except Exception as e:
        answer = f"Agent error: {e}"
    finally:
        if tmp_path:
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

    return JSONResponse({"answer": answer})


# ── Course submission endpoint (NEW) ────────────────────────────────────────

@app.post("/run-and-submit")
async def run_and_submit():
    """
    1. Fetch the 20 course questions from the HF course API.
    2. Run the agent on each question.
    3. POST answers to the course leaderboard.
    4. Return the score and leaderboard response.
    """
    if not HF_TOKEN:
        return JSONResponse(
            {"error": "HF_TOKEN secret not set on this Space."},
            status_code=500
        )

    # ── Step 1: Fetch the 20 course questions ───────────────────────────────
    log.info("Fetching course questions from %s/questions", COURSE_API)
    try:
        resp = requests.get(f"{COURSE_API}/questions", timeout=30)
        resp.raise_for_status()
        questions = resp.json()
    except Exception as e:
        log.error("Failed to fetch questions: %s", e)
        return JSONResponse(
            {"error": f"Failed to fetch questions from course API: {e}"},
            status_code=500
        )

    log.info("Fetched %d questions", len(questions))

    # ── Step 2: Run agent on each question ──────────────────────────────────
    answers = []
    agent = build_agent(HF_TOKEN)

    for i, q in enumerate(questions):
        task_id  = q.get("task_id", f"task_{i}")
        question_text = q.get("Question", q.get("question", ""))

        # Handle optional attached file
        fname = (q.get("file_name") or "").strip()
        full_q = question_text.strip()

        if fname:
            try:
                from huggingface_hub import hf_hub_download
                attached = hf_hub_download(
                    repo_id="gaia-benchmark/GAIA",
                    filename=f"2023/validation/{fname}",
                    repo_type="dataset",
                    token=HF_TOKEN,
                )
                full_q += f"\n\nAttached file path: {attached}"
                log.info("[%d/%d] Downloaded attachment: %s", i+1, len(questions), fname)
            except Exception as fe:
                log.warning("[%d/%d] File download failed (%s): %s", i+1, len(questions), fname, fe)

        log.info("[%d/%d] Running agent on: %s", i+1, len(questions), question_text[:80])

        try:
            answer = str(agent.run(full_q))
        except Exception as e:
            answer = f"ERROR: {e}"
            log.error("[%d/%d] Agent error: %s", i+1, len(questions), e)

        log.info("[%d/%d] Answer: %s", i+1, len(questions), answer[:80])

        answers.append({
            "task_id": task_id,
            "submitted_answer": answer,
        })

    # ── Step 3: POST to course leaderboard ──────────────────────────────────
    log.info("Submitting %d answers to course leaderboard", len(answers))
    payload = {
        "username":   HF_USERNAME,
        "agent_code": SPACE_URL,
        "answers":    answers,
    }

    try:
        submit_resp = requests.post(
            f"{COURSE_API}/submit",
            json=payload,
            timeout=60,
        )
        submit_resp.raise_for_status()
        result = submit_resp.json()
    except Exception as e:
        log.error("Submission failed: %s", e)
        # Return partial results so answers are not lost
        return JSONResponse({
            "error": f"Submission to course API failed: {e}",
            "answers_generated": len(answers),
            "answers": answers,
        }, status_code=500)

    log.info("Submission result: %s", result)

    # Return score + full answer list for debugging
    return JSONResponse({
        "submission_result": result,
        "answers_submitted": len(answers),
        "answers": answers,
    })


# ── Health check ─────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {"status": "ok", "hf_token_set": bool(HF_TOKEN)}


# ── Local dev entrypoint ──────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)