Blame for app.py

"""
GAIA Unit 4 final-assignment agent (GPT-4o, tool-calling loop).

Keeps the official course workflow: Gradio + Hugging Face OAuth login,
GET /questions, POST /submit with {username, agent_code, answers}.

Secrets (Space Settings -> Variables and secrets):
  OPENAI_API_KEY  required
  HF_TOKEN        optional but recommended: lets the agent fetch task
                  attachments straight from the gated GAIA dataset if the
                  course /files/{task_id} endpoint fails (accept the dataset
                  conditions at huggingface.co/datasets/gaia-benchmark/GAIA).
"""
from __future__ import annotations

import base64
import io
import json
import mimetypes
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

import gradio as gr
import pandas as pd
import requests
from openai import OpenAI

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
SCORING_API = os.getenv("SCORING_API", "https://agents-course-unit4-scoring.hf.space")
GAIA_DATASET_ID = "gaia-benchmark/GAIA"
GAIA_FILE_DIRS = ("2023/validation", "2023/test")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_AUDIO_MODEL = os.getenv("OPENAI_AUDIO_MODEL", "whisper-1")
HF_TOKEN = os.getenv("HF_TOKEN", "")

MAX_STEPS = int(os.getenv("GAIA_MAX_STEPS", "9"))        # tool-loop iterations per question
MAX_COMPLETION_TOKENS = 900
HTTP_TIMEOUT = 30
FETCH_CHARS = 6000          # chars returned per fetch_url call
SEARCH_RESULTS = 6
PY_TIMEOUT = 40
PY_OUTPUT_CHARS = 4000
WHISPER_MAX_BYTES = 24 * 1024 * 1024
BROWSER_UA = (
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/124.0 Safari/537.36"
)

TEXT_EXT = {".txt", ".md", ".json", ".jsonl", ".csv", ".tsv", ".xml", ".html",
            ".htm", ".log", ".yaml", ".yml", ".jsonld"}
IMAGE_EXT = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
AUDIO_EXT = {".mp3", ".wav", ".m4a", ".mpga", ".mpeg", ".ogg", ".flac", ".webm", ".mp4"}


# Optional ZeroGPU startup compatibility. Inference is remote (OpenAI), so this
# function is never called. It only prevents "No @spaces.GPU function detected"
# if the Space hardware is still set to ZeroGPU. Harmless on CPU Basic.
try:
    import spaces  # type: ignore

    @spaces.GPU(duration=1)
    def _zerogpu_startup_probe():
        return None
except Exception:
    pass


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def truncate(text: str, limit: int) -> str:
    text = text or ""
    if len(text) <= limit:
        return text
    return text[:limit] + f"\n...[truncated, {len(text) - limit} more chars]"


def clean_answer(value: Any) -> str:
    """Extract and normalise the final answer for EXACT-MATCH scoring."""
    text = str(value or "").strip()
    matches = list(re.finditer(r"final answer\s*[:\-]\s*", text, flags=re.IGNORECASE))
    if matches:
        text = text[matches[-1].end():]
    text = text.strip().splitlines()[0].strip() if text.strip() else ""
    if len(text) >= 2 and text[0] == text[-1] and text[0] in {"'", '"'}:
        text = text[1:-1].strip()
    text = text.replace("**", "").strip()
    # "1,234" -> "1234" (GAIA numbers carry no thousands separators); lists use ", "
    if re.fullmatch(r"-?\d{1,3}(,\d{3})+(\.\d+)?", text):
        text = text.replace(",", "")
    # drop a trailing full stop on short answers
    if text.endswith(".") and len(text.split()) <= 6 and not re.search(r"\b[A-Z]\.$", text):
        text = text[:-1].strip()
    return text


def image_to_data_url(path: str, max_side: int = 1568) -> str:
    """Downscale large images (cost control) and return a data URL."""
    try:
        from PIL import Image
        with Image.open(path) as img:
            img.load()
            if img.mode not in ("RGB", "L"):
                img = img.convert("RGB")
            img.thumbnail((max_side, max_side))
            buf = io.BytesIO()
            img.save(buf, format="PNG" if img.mode == "L" else "JPEG", quality=90) \
                if img.mode != "L" else img.save(buf, format="PNG")
            mime = "image/png" if img.mode == "L" else "image/jpeg"
            return f"data:{mime};base64,{base64.b64encode(buf.getvalue()).decode()}"
    except Exception:
        mime = mimetypes.guess_type(path)[0] or "image/png"
        return f"data:{mime};base64,{base64.b64encode(Path(path).read_bytes()).decode()}"


def html_to_text(html: str) -> str:
    """HTML -> readable text; tables become 'a | b | c' rows."""
    from bs4 import BeautifulSoup
    soup = BeautifulSoup(html, "lxml") if _has_lxml() else BeautifulSoup(html, "html.parser")
    for tag in soup(["script", "style", "noscript", "svg", "nav", "footer", "form", "aside"]):
        tag.decompose()
    for sup in soup.select("sup.reference, .mw-editsection, .navbox, .reflist"):
        sup.decompose()
    for table in soup.find_all("table"):
        rows = []
        for tr in table.find_all("tr"):
            cells = [c.get_text(" ", strip=True) for c in tr.find_all(["th", "td"])]
            if any(cells):
                rows.append(" | ".join(cells))
        table.replace_with(soup.new_string("\n[TABLE]\n" + "\n".join(rows) + "\n[/TABLE]\n"))
    text = soup.get_text("\n")
    text = re.sub(r"[ \t\xa0]+", " ", text)
    text = re.sub(r"\n\s*\n+", "\n", text)
    return text.strip()


def _has_lxml() -> bool:
    try:
        import lxml  # noqa: F401
        return True
    except Exception:
        return False


def youtube_id(url: str) -> str | None:
    m = re.search(r"(?:v=|youtu\.be/|/shorts/|/embed/)([A-Za-z0-9_-]{11})", url or "")
    return m.group(1) if m else None


def parse_vtt(raw: str) -> str:
    lines, seen = [], set()
    for line in raw.splitlines():
        line = line.strip()
        if (not line or line == "WEBVTT" or "-->" in line or line.isdigit()
                or line.startswith(("NOTE", "Kind:", "Language:"))):
            continue
        line = re.sub(r"<[^>]+>", "", line)
        line = line.replace("&amp;", "&").replace("&quot;", '"').replace("&#39;", "'")
        if line and line not in seen:
            seen.add(line)
            lines.append(line)
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Attachment retrieval
# ---------------------------------------------------------------------------
def download_attachment(task_id: str, file_name: str, dest_dir: str) -> str | None:
    """Course API first, then the gated GAIA dataset (needs HF_TOKEN)."""
    if not file_name:
        return None
    dest = Path(dest_dir) / Path(file_name).name

    # 1) course scoring API
    try:
        r = requests.get(f"{SCORING_API}/files/{task_id}", timeout=HTTP_TIMEOUT)
        ctype = r.headers.get("content-type", "")
        looks_like_error = "json" in ctype and r.content[:1] in (b"{", b"[") and len(r.content) < 600
        if r.ok and r.content and not looks_like_error:
            dest.write_bytes(r.content)
            print(f"[files] {task_id}: course API ok ({len(r.content)} bytes)")
            return str(dest)
        print(f"[files] {task_id}: course API returned {r.status_code}; trying GAIA dataset")
    except requests.RequestException as exc:
        print(f"[files] {task_id}: course API error: {exc}")

    # 2) GAIA dataset: the file lives at <split dir>/<file_name>
    try:
        from huggingface_hub import hf_hub_download
        for folder in GAIA_FILE_DIRS:
            try:
                local = hf_hub_download(
                    repo_id=GAIA_DATASET_ID,
                    filename=f"{folder}/{Path(file_name).name}",
                    repo_type="dataset",
                    token=HF_TOKEN or None,
                )
                shutil.copy(local, dest)
                print(f"[files] {task_id}: fetched from GAIA dataset ({folder})")
                return str(dest)
            except Exception as exc:  # try next split
                last = exc
        print(f"[files] {task_id}: GAIA dataset download failed: {last}")
    except Exception as exc:
        print(f"[files] {task_id}: huggingface_hub unavailable: {exc}")
    return None


# ---------------------------------------------------------------------------
# Agent
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """You are an expert research agent solving GAIA benchmark questions. \
Answers are graded by EXACT MATCH, so correctness and format both matter.

How to work:
- Plan briefly, then use tools. Never answer facts from memory when they can be verified: use web_search, then fetch_url to read the actual page (snippets are not enough). Prefer Wikipedia, official sites, and primary sources. Use fetch_url's `find` argument to jump to relevant passages on long pages.
- Do ALL arithmetic, counting, sorting, filtering, date maths and table/spreadsheet analysis with run_python, never by eye. Attachments are saved locally; their path is given in the task. For Excel use pandas/openpyxl (openpyxl for cell colours/formatting).
- Python attachments: run them with run_python (e.g. subprocess or exec of the file) to get the real output rather than guessing.
- Questions with a YouTube link: call youtube_transcript first; for questions about what is visible on screen call youtube_frames. Audio attachments are already transcribed for you.
- If the question text looks reversed or encoded, decode it first (run_python can do this).
- Read constraints carefully (dates, 'as of', 'between', 'at least', which edition/version, rounding, units requested).
- If a source fails, try another query or source. Do not give up early, but do not loop endlessly; stop researching once you have verified evidence.
- Never claim to have viewed a file/page/video you could not access. If something is truly unobtainable, give your best-supported answer.

Final answer format (very important):
- End with a line: FINAL ANSWER: <answer>
- Number -> digits only, no commas, no units or % or $ unless the question demands them.
- String -> as few words as possible, no articles (a/an/the), no abbreviations unless asked, digits in plain form.
- List -> comma-separated, applying the above per element, in the order requested.
- No explanations, no markdown, nothing after the final answer line."""


TOOLS = [
    {"type": "function", "function": {
        "name": "web_search",
        "description": "Search the web (DuckDuckGo, Wikipedia fallback). Returns titles, URLs and snippets. Follow up with fetch_url on promising results.",
        "parameters": {"type": "object", "properties": {
            "query": {"type": "string"}},
            "required": ["query"]}}},
    {"type": "function", "function": {
        "name": "fetch_url",
        "description": "Download a web page / PDF / text file and return readable text (tables as 'a | b' rows). Works for Wikipedia and web.archive.org URLs. Use `find` (keywords, case-insensitive) to get only matching passages from long pages, or `offset` to page through.",
        "parameters": {"type": "object", "properties": {
            "url": {"type": "string"},
            "find": {"type": "string", "description": "Optional keyword/regex; returns passages around matches."},
            "offset": {"type": "integer", "description": "Character offset to start from (default 0)."}},
            "required": ["url"]}}},
    {"type": "function", "function": {
        "name": "run_python",
        "description": "Execute Python 3 (pandas, openpyxl, numpy, requests available) in the task's working directory and return stdout/stderr. Always print() results. Use it for all calculation and for reading/processing attachments.",
        "parameters": {"type": "object", "properties": {
            "code": {"type": "string"}},
            "required": ["code"]}}},
    {"type": "function", "function": {
        "name": "youtube_transcript",
        "description": "Get the caption transcript (or, if none, a Whisper audio transcript) of a YouTube video.",
        "parameters": {"type": "object", "properties": {
            "url": {"type": "string"}},
            "required": ["url"]}}},
    {"type": "function", "function": {
        "name": "youtube_frames",
        "description": "Sample frames from a YouTube video and show them to you as images. Use for questions about what is visible. Keep num_frames small.",
        "parameters": {"type": "object", "properties": {
            "url": {"type": "string"},
            "num_frames": {"type": "integer", "description": "1-12, default 8"},
            "start_seconds": {"type": "number"},
            "end_seconds": {"type": "number"}},
            "required": ["url"]}}},
    {"type": "function", "function": {
        "name": "view_image",
        "description": "Show a local image file or an image URL to you for visual inspection.",
        "parameters": {"type": "object", "properties": {
            "path_or_url": {"type": "string"}},
            "required": ["path_or_url"]}}},
    {"type": "function", "function": {
        "name": "transcribe_audio",
        "description": "Transcribe a local audio file (mp3/wav/m4a/...) with Whisper.",
        "parameters": {"type": "object", "properties": {
            "path": {"type": "string"}},
            "required": ["path"]}}},
]


class GaiaAgent:
    def __init__(self):
        if not OPENAI_API_KEY:
            raise ValueError(
                "Missing OPENAI_API_KEY. Add it under Space Settings -> Variables and secrets."
            )
        self.client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL,
                             timeout=120, max_retries=4)
        self.http = requests.Session()
        self.http.headers.update({"User-Agent": BROWSER_UA, "Accept-Language": "en-US,en;q=0.9"})
        self.workdir = ""
        self.pending_images: list[dict] = []
        self.usage = {"prompt": 0, "completion": 0}
        print(f"Initialized agent with model {OPENAI_MODEL}")

    # ------------------------------ tools ---------------------------------
    def web_search(self, query: str) -> str:
        results = []
        try:
            from ddgs import DDGS
            for attempt in range(2):
                try:
                    results = list(DDGS().text(query, max_results=SEARCH_RESULTS) or [])
                    if results:
                        break
                except Exception as exc:
                    print(f"[search] ddgs attempt {attempt + 1} failed: {exc}")
                    time.sleep(1.5)
        except Exception as exc:
            print(f"[search] ddgs unavailable: {exc}")

        lines = []
        for r in results:
            lines.append(f"- {r.get('title', '')}\n  URL: {r.get('href') or r.get('url', '')}\n"
                         f"  {truncate(r.get('body', ''), 300)}")

        # Wikipedia fallback / supplement (always cheap and reliable)
        if len(results) < 3:
            try:
                resp = self.http.get(
                    "https://en.wikipedia.org/w/api.php",
                    params={"action": "query", "list": "search", "srsearch": query,
                            "format": "json", "srlimit": 5},
                    timeout=HTTP_TIMEOUT)
                for item in resp.json().get("query", {}).get("search", []):
                    title = item["title"]
                    snippet = re.sub(r"<[^>]+>", "", item.get("snippet", ""))
                    url = "https://en.wikipedia.org/wiki/" + title.replace(" ", "_")
                    lines.append(f"- [Wikipedia] {title}\n  URL: {url}\n  {snippet}")
            except Exception as exc:
                print(f"[search] wikipedia fallback failed: {exc}")

        return "\n".join(lines) if lines else "No results. Try a different query."

    def fetch_url(self, url: str, find: str | None = None, offset: int = 0) -> str:
        url = (url or "").strip()
        if not url.startswith("http"):
            return "Invalid URL."
        try:
            resp = self.http.get(url, timeout=HTTP_TIMEOUT)
            resp.raise_for_status()
        except Exception as exc:
            return f"Fetch failed: {exc}"

        ctype = resp.headers.get("content-type", "").lower()
        try:
            if "pdf" in ctype or url.lower().endswith(".pdf"):
                from pypdf import PdfReader
                reader = PdfReader(io.BytesIO(resp.content))
                text = "\n".join((p.extract_text() or "") for p in reader.pages[:60])
            elif "html" in ctype or "<html" in resp.text[:500].lower():
                text = html_to_text(resp.text)
            elif any(t in ctype for t in ("text", "json", "xml", "csv")):
                text = resp.text
            else:
                saved = Path(self.workdir or tempfile.gettempdir()) / (Path(url.split("?")[0]).name or "download.bin")
                saved.write_bytes(resp.content)
                return f"Binary content ({ctype}) saved to {saved}. Process it with run_python."
        except Exception as exc:
            return f"Could not parse page: {exc}"

        if find:
            try:
                pattern = re.compile(find, re.IGNORECASE)
            except re.error:
                pattern = re.compile(re.escape(find), re.IGNORECASE)
            spans, last_end = [], -1
            for m in pattern.finditer(text):
                if m.start() < last_end:
                    continue
                a, b = max(0, m.start() - 350), min(len(text), m.end() + 650)
                spans.append(text[a:b])
                last_end = b
                if len(spans) >= 8:
                    break
            if spans:
                return truncate("\n---\n".join(spans), FETCH_CHARS) + \
                    f"\n[page length {len(text)} chars]"
            return f"No match for '{find}'. Page length {len(text)} chars; first part:\n" + \
                truncate(text[:2000], 2000)

        chunk = text[offset: offset + FETCH_CHARS]
        more = len(text) - (offset + FETCH_CHARS)
        suffix = f"\n[page length {len(text)} chars; next offset {offset + FETCH_CHARS}]" if more > 0 else ""
        return chunk + suffix

    def run_python(self, code: str) -> str:
        env = {k: v for k, v in os.environ.items()
               if k not in {"OPENAI_API_KEY", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}}
        env["MPLBACKEND"] = "Agg"
        script = Path(self.workdir) / "_snippet.py"
        script.write_text(code, encoding="utf-8")
        try:
            proc = subprocess.run(
                [sys.executable, str(script)], cwd=self.workdir, env=env,
                capture_output=True, text=True, timeout=PY_TIMEOUT)
            out = proc.stdout
            if proc.stderr:
                out += ("\n[stderr]\n" + proc.stderr[-1500:])
            return truncate(out.strip() or "(no output - remember to print())", PY_OUTPUT_CHARS)
        except subprocess.TimeoutExpired:
            return f"Execution timed out after {PY_TIMEOUT}s."
        except Exception as exc:
            return f"Execution error: {exc}"

    def transcribe_audio(self, path: str) -> str:
        p = Path(path)
        if not p.exists():
            return f"File not found: {path}"
        if p.stat().st_size > WHISPER_MAX_BYTES:
            return "Audio file larger than 24MB; cannot transcribe."
        try:
            with open(p, "rb") as fh:
                resp = self.client.audio.transcriptions.create(model=OPENAI_AUDIO_MODEL, file=fh)
            return resp.text or "(empty transcript)"
        except Exception as exc:
            return f"Transcription failed: {exc}"

    def view_image(self, path_or_url: str) -> str:
        try:
            if path_or_url.startswith("http"):
                resp = self.http.get(path_or_url, timeout=HTTP_TIMEOUT)
                resp.raise_for_status()
                tmp = Path(self.workdir) / "_remote_image"
                tmp.write_bytes(resp.content)
                path = str(tmp)
            else:
                path = path_or_url
            self.pending_images.append(
                {"type": "image_url", "image_url": {"url": image_to_data_url(path), "detail": "high"}})
            return "Image loaded; it is shown in the next message."
        except Exception as exc:
            return f"Could not load image: {exc}"

    # ---- YouTube ----
    def youtube_transcript(self, url: str) -> str:
        vid = youtube_id(url)
        if not vid:
            return "Not a valid YouTube URL."
        errors = []

        # 1) youtube-transcript-api (handles old and new API shapes)
        try:
            from youtube_transcript_api import YouTubeTranscriptApi
            segments = None
            try:
                api = YouTubeTranscriptApi()
                fetched = api.fetch(vid, languages=["en", "en-US", "en-GB"])
                segments = [getattr(s, "text", None) or s["text"] for s in fetched]
            except AttributeError:
                data = YouTubeTranscriptApi.get_transcript(vid, languages=["en", "en-US", "en-GB"])
                segments = [s["text"] for s in data]
            if segments:
                return truncate("TRANSCRIPT:\n" + " ".join(segments), 12000)
        except Exception as exc:
            errors.append(f"transcript-api: {str(exc)[:200]}")

        # 2) yt-dlp subtitles
        try:
            import yt_dlp
            with tempfile.TemporaryDirectory() as tmp:
                opts = {"skip_download": True, "writesubtitles": True, "writeautomaticsub": True,
                        "subtitleslangs": ["en", "en-US", "en-GB", "en.*"], "subtitlesformat": "vtt",
                        "outtmpl": str(Path(tmp) / "yt.%(ext)s"), "quiet": True,
                        "no_warnings": True, "noplaylist": True, "socket_timeout": 20}
                with yt_dlp.YoutubeDL(opts) as ydl:
                    info = ydl.extract_info(url, download=True)
                files = list(Path(tmp).glob("*.vtt"))
                if files:
                    text = parse_vtt(files[0].read_text(encoding="utf-8", errors="replace"))
                    if text:
                        return truncate(f"TITLE: {(info or {}).get('title', '')}\nTRANSCRIPT:\n{text}", 12000)
        except Exception as exc:
            errors.append(f"yt-dlp subtitles: {str(exc)[:200]}")

        # 3) audio -> Whisper (short videos only)
        try:
            import yt_dlp
            with tempfile.TemporaryDirectory() as tmp:
                opts = {"format": "bestaudio[filesize<24M]/worstaudio/worst",
                        "outtmpl": str(Path(tmp) / "audio.%(ext)s"), "quiet": True,
                        "no_warnings": True, "noplaylist": True, "socket_timeout": 20,
                        "max_filesize": WHISPER_MAX_BYTES}
                with yt_dlp.YoutubeDL(opts) as ydl:
                    info = ydl.extract_info(url, download=True)
                if (info or {}).get("duration", 0) and info["duration"] > 1500:
                    errors.append("video longer than 25 minutes; audio transcription skipped")
                else:
                    files = [p for p in Path(tmp).glob("audio.*")]
                    if files:
                        text = self.transcribe_audio(str(files[0]))
                        return truncate("WHISPER TRANSCRIPT (no captions available):\n" + text, 12000)
        except Exception as exc:
            errors.append(f"audio download: {str(exc)[:200]}")

        return ("Could not retrieve transcript. Reasons: " + " | ".join(errors) +
                ". Try web_search for a transcript/summary of this video instead.")

    def youtube_frames(self, url: str, num_frames: int = 8,
                       start_seconds: float | None = None, end_seconds: float | None = None) -> str:
        num_frames = max(1, min(int(num_frames or 8), 12))
        try:
            import cv2
            import yt_dlp
            from PIL import Image
        except Exception as exc:
            return f"Frame extraction dependencies missing: {exc}"
        try:
            with tempfile.TemporaryDirectory() as tmp:
                opts = {"format": "worst[ext=mp4]/worst[height>=240]/worst",
                        "outtmpl": str(Path(tmp) / "video.%(ext)s"), "quiet": True,
                        "no_warnings": True, "noplaylist": True, "socket_timeout": 25,
                        "max_filesize": 120 * 1024 * 1024}
                with yt_dlp.YoutubeDL(opts) as ydl:
                    ydl.download([url])
                vids = [p for p in Path(tmp).glob("video.*") if p.suffix.lower() in {".mp4", ".mkv", ".webm"}]
                if not vids:
                    return "Video download failed (no file produced)."
                cap = cv2.VideoCapture(str(vids[0]))
                fps = cap.get(cv2.CAP_PROP_FPS) or 25
                total = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0
                duration = total / fps if fps else 0
                if duration <= 0:
                    return "Could not read video duration."
                t0 = max(0.0, start_seconds or 0.0)
                t1 = min(duration, end_seconds if end_seconds else duration)
                added = 0
                for i in range(num_frames):
                    t = t0 + (t1 - t0) * (i + 0.5) / num_frames
                    cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
                    ok, frame = cap.read()
                    if not ok:
                        continue
                    img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                    img.thumbnail((768, 768))
                    buf = io.BytesIO()
                    img.save(buf, format="JPEG", quality=80)
                    self.pending_images.append({"type": "image_url", "image_url": {
                        "url": "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode(),
                        "detail": "auto"}})
                    added += 1
                cap.release()
                if not added:
                    return "No frames could be decoded."
                return (f"{added} frames sampled between {t0:.0f}s and {t1:.0f}s "
                        f"(video length {duration:.0f}s); shown in the next message.")
        except Exception as exc:
            return f"Frame extraction failed: {str(exc)[:300]}"

    def dispatch(self, name: str, args: dict) -> str:
        try:
            if name == "web_search":
                return self.web_search(args.get("query", ""))
            if name == "fetch_url":
                return self.fetch_url(args.get("url", ""), args.get("find"), int(args.get("offset") or 0))
            if name == "run_python":
                return self.run_python(args.get("code", ""))
            if name == "youtube_transcript":
                return self.youtube_transcript(args.get("url", ""))
            if name == "youtube_frames":
                return self.youtube_frames(args.get("url", ""), args.get("num_frames", 8),
                                           args.get("start_seconds"), args.get("end_seconds"))
            if name == "view_image":
                return self.view_image(args.get("path_or_url", ""))
            if name == "transcribe_audio":
                return self.transcribe_audio(args.get("path", ""))
            return f"Unknown tool: {name}"
        except Exception as exc:
            return f"Tool {name} crashed: {exc}"

    # ------------------------- attachment preview --------------------------
    def prepare_attachment(self, task_id: str, file_name: str) -> tuple[str, list[dict]]:
        """Return (description for the prompt, images to attach)."""
        if not file_name:
            return "", []
        path = download_attachment(task_id, file_name, self.workdir)
        if not path:
            return (f"ATTACHMENT: '{file_name}' exists for this task but could NOT be downloaded "
                    "(course API and GAIA dataset both failed). Say so honestly if the answer "
                    "depends on it; otherwise answer from other evidence."), []

        ext = Path(path).suffix.lower()
        header = f"ATTACHMENT saved at: {path} (type {ext or 'unknown'}, {Path(path).stat().st_size} bytes)\n"
        images: list[dict] = []
        try:
            if ext in IMAGE_EXT:
                images.append({"type": "image_url",
                               "image_url": {"url": image_to_data_url(path), "detail": "high"}})
                return header + "The image is attached to this message.", images
            if ext in {".mp3", ".wav", ".m4a", ".mpga", ".mpeg", ".ogg", ".flac"}:
                return header + "AUDIO TRANSCRIPT (Whisper):\n" + self.transcribe_audio(path), images
            if ext in {".xlsx", ".xlsm", ".xls"}:
                sheets = pd.read_excel(path, sheet_name=None, header=None)
                parts = []
                for name, frame in sheets.items():
                    parts.append(f"SHEET '{name}' shape={frame.shape}\n{truncate(frame.head(25).to_csv(index=False, header=False), 2500)}")
                return header + "Spreadsheet preview (first rows, no header inference). " \
                    "Use run_python with pandas/openpyxl for the full data:\n" + "\n".join(parts), images
            if ext == ".py":
                code = Path(path).read_text(encoding="utf-8", errors="replace")
                return header + "PYTHON SOURCE (execute it with run_python to get real output):\n" + truncate(code, 5000), images
            if ext in {".csv", ".tsv"}:
                return header + "File head:\n" + truncate(Path(path).read_text(encoding="utf-8", errors="replace"), 2500), images
            if ext in TEXT_EXT:
                return header + "CONTENT:\n" + truncate(Path(path).read_text(encoding="utf-8", errors="replace"), 6000), images
            if ext == ".pdf":
                from pypdf import PdfReader
                text = "\n".join((p.extract_text() or "") for p in PdfReader(path).pages[:30])
                return header + "PDF TEXT:\n" + truncate(text, 6000), images
            if ext == ".docx":
                from docx import Document
                doc = Document(path)
                lines = [p.text for p in doc.paragraphs if p.text.strip()]
                for t in doc.tables:
                    for row in t.rows:
                        lines.append(" | ".join(c.text for c in row.cells))
                return header + "DOCX TEXT:\n" + truncate("\n".join(lines), 6000), images
            if ext == ".pptx":
                from pptx import Presentation
                out = []
                for i, slide in enumerate(Presentation(path).slides, 1):
                    texts = [s.text for s in slide.shapes if getattr(s, "has_text_frame", False)]
                    out.append(f"SLIDE {i}: " + " / ".join(t for t in texts if t))
                return header + truncate("\n".join(out), 6000), images
            if ext == ".zip":
                import zipfile
                with zipfile.ZipFile(path) as z:
                    z.extractall(self.workdir)
                    return header + "ZIP extracted into the working directory. Files: " + ", ".join(z.namelist()[:40]), images
            if ext in {".mp4", ".mov", ".mkv"}:
                return header + "Video file; use run_python (cv2) to inspect frames if needed.", images
        except Exception as exc:
            return header + f"(Preview failed: {exc}. Inspect it with run_python.)", images
        return header + "Unknown format; inspect it with run_python.", images

    # ------------------------------ main loop ------------------------------
    def _chat(self, messages: list[dict], allow_tools: bool = True):
        kwargs: dict[str, Any] = dict(model=OPENAI_MODEL, temperature=0,
                                      max_tokens=MAX_COMPLETION_TOKENS, messages=messages)
        if allow_tools:
            kwargs.update(tools=TOOLS, tool_choice="auto")
        resp = self.client.chat.completions.create(**kwargs)
        if resp.usage:
            self.usage["prompt"] += resp.usage.prompt_tokens
            self.usage["completion"] += resp.usage.completion_tokens
        return resp.choices[0].message

    def __call__(self, task: dict) -> str:
        task_id = str(task.get("task_id", ""))
        question = str(task.get("question", ""))
        file_name = str(task.get("file_name") or "")
        print(f"\n=== Task {task_id} | attachment={file_name or 'none'}")

        self.workdir = tempfile.mkdtemp(prefix=f"gaia_{task_id[:8]}_")
        self.pending_images = []
        try:
            attachment_text, attachment_images = self.prepare_attachment(task_id, file_name)
            text = f"QUESTION:\n{question}"
            if attachment_text:
                text += f"\n\n{attachment_text}"
            user_content: Any = [{"type": "text", "text": text}] + attachment_images \
                if attachment_images else text

            messages: list[dict] = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_content},
            ]

            final_text = ""
            for step in range(MAX_STEPS):
                msg = self._chat(messages)
                if not msg.tool_calls:
                    final_text = msg.content or ""
                    break
                messages.append({
                    "role": "assistant", "content": msg.content or "",
                    "tool_calls": [{"id": tc.id, "type": "function",
                                    "function": {"name": tc.function.name,
                                                 "arguments": tc.function.arguments}}
                                   for tc in msg.tool_calls]})
                for tc in msg.tool_calls:
                    try:
                        args = json.loads(tc.function.arguments or "{}")
                    except json.JSONDecodeError:
                        args = {}
                    print(f"[step {step + 1}] {tc.function.name}({truncate(json.dumps(args), 160)})")
                    result = self.dispatch(tc.function.name, args)
                    messages.append({"role": "tool", "tool_call_id": tc.id,
                                     "content": truncate(result, 12500)})
                if self.pending_images:
                    messages.append({"role": "user", "content": [
                        {"type": "text", "text": "Images requested by your tool call:"},
                        *self.pending_images[:12]]})
                    self.pending_images = []
            else:
                # Out of steps: force a final answer from what has been gathered.
                messages.append({"role": "user", "content":
                                 "Step limit reached. Give your best-supported answer now as "
                                 "'FINAL ANSWER: ...' with no further tool use."})
                final_text = self._chat(messages, allow_tools=False).content or ""

            answer = clean_answer(final_text)
            print(f"Answer [{task_id}]: {answer}  | tokens so far: {self.usage}")
            return answer
        finally:
            shutil.rmtree(self.workdir, ignore_errors=True)


# ---------------------------------------------------------------------------
# Evaluation and submission (official course workflow)
# ---------------------------------------------------------------------------
def run_and_submit_all(profile: gr.OAuthProfile | None, progress=gr.Progress()):
    if not profile:
        return "Please log in to Hugging Face with the button.", None
    username = str(profile.username).strip()
    print(f"User logged in: {username}")

    space_id = os.getenv("SPACE_ID", "")
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else ""

    try:
        agent = GaiaAgent()
    except Exception as exc:
        return f"Error initializing agent: {exc}", None

    try:
        resp = requests.get(f"{SCORING_API}/questions", timeout=HTTP_TIMEOUT)
        resp.raise_for_status()
        questions = resp.json()
        if not isinstance(questions, list) or not questions:
            return "Fetched questions list is empty or invalid.", None
        print(f"Fetched {len(questions)} questions.")
    except Exception as exc:
        return f"Error fetching questions: {exc}", None

    results_log, answers_payload = [], []
    total = len(questions)
    for index, item in enumerate(questions, start=1):
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            continue
        progress((index - 1) / total, desc=f"Solving question {index}/{total}")
        try:
            answer = agent(item) or "unknown"
        except Exception as exc:
            print(f"Task {task_id} failed: {exc}")
            answer = "unknown"
        answers_payload.append({"task_id": task_id, "submitted_answer": answer})
        results_log.append({"Task ID": task_id, "Question": question_text,
                            "Attachment": item.get("file_name", ""),
                            "Submitted Answer": answer})

    results_df = pd.DataFrame(results_log)
    if not answers_payload:
        return "Agent did not produce any answers to submit.", results_df

    print(f"Total tokens used: {agent.usage}")
    submission = {"username": username, "agent_code": agent_code, "answers": answers_payload}
    progress(0.97, desc="Submitting answers")
    try:
        resp = requests.post(f"{SCORING_API}/submit", json=submission, timeout=120)
        resp.raise_for_status()
        data = resp.json()
        status = (
            "Submission Successful!\n"
            f"User: {data.get('username')}\n"
            f"Overall Score: {data.get('score', 'N/A')}% "
            f"({data.get('correct_count', '?')}/{data.get('total_attempted', '?')} correct)\n"
            f"Message: {data.get('message', 'No message received.')}\n"
            f"OpenAI tokens this run: {agent.usage['prompt']} in / {agent.usage['completion']} out"
        )
        return status, results_df
    except requests.exceptions.HTTPError as exc:
        detail = exc.response.text[:500] if exc.response is not None else ""
        return f"Submission Failed: HTTP {getattr(exc.response, 'status_code', '?')}: {detail}", results_df
    except Exception as exc:
        return f"Submission Failed: {exc}", results_df


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Final Assignment - GPT-4o Tool-Calling Agent")
    gr.Markdown(
        """
**Instructions**
1. Add `OPENAI_API_KEY` (required) and `HF_TOKEN` (recommended, for GAIA attachment fallback) as Space secrets.
2. Log in with your Hugging Face account below.
3. Click **Run Evaluation & Submit All Answers**. The run takes several minutes.

The agent searches the web, reads pages, runs Python, and handles audio, images, YouTube and spreadsheets.
Keep this Space public so `agent_code` can be verified.
"""
    )
    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers", variant="primary")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=7, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)
    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])


if __name__ == "__main__":
    print(f"Starting GAIA agent. Model: {OPENAI_MODEL}. Space: {os.getenv('SPACE_ID', 'local')}")
    demo.launch(debug=True, share=False)