api/client.py

from __future__ import annotations

import json
import shutil
import time
from pathlib import Path
from typing import Any

import requests
from datasets import load_dataset
from huggingface_hub import snapshot_download

from api.schemas import GaiaQuestion
from config import Settings


class GaiaApiClient:
    """Client for the Hugging Face AI Agents course evaluation API."""

    def __init__(
        self,
        settings: Settings,
        session: requests.Session | None = None,
    ) -> None:
        self.settings = settings
        self.session = session or requests.Session()

        self.session.headers.update(
            {
                "User-Agent": "Vertex-Agent/1.0",
                "Accept": "application/json",
            }
        )

    def _request(
        self,
        method: str,
        path: str,
        **kwargs: Any,
    ) -> requests.Response:
        """Send a request to the course API."""

        base_url = self.settings.course_api_url.rstrip("/")
        url = f"{base_url}{path}"

        response = self.session.request(
            method=method,
            url=url,
            timeout=self.settings.request_timeout,
            **kwargs,
        )

        response.raise_for_status()
        return response

    def get_questions(self) -> list[GaiaQuestion]:
        """Load local debug questions if available, otherwise use the course API."""

        debug_file = Path("questions_debug.json")

        if debug_file.exists():
            print(f"Using local questions: {debug_file.resolve()}")

            with debug_file.open(
                "r",
                encoding="utf-8",
            ) as f:
                data = json.load(f)

            if not isinstance(data, list):
                raise RuntimeError(
                    "questions_debug.json must contain a JSON list."
                )

            return [
                GaiaQuestion.from_api(item)
                for item in data
            ]

        print("Using course API questions...")

        response = self._request(
            method="GET",
            path="/questions",
        )

        data = response.json()

        if not isinstance(data, list):
            raise RuntimeError(
                "The /questions endpoint did not return a list."
            )

        return [
            GaiaQuestion.from_api(item)
            for item in data
        ]

    def download_file(
        self,
        question: GaiaQuestion,
    ) -> Path | None:
        """
        Download a question attachment.

        The course API is attempted first. If that endpoint does not
        provide the attachment, the official gated GAIA dataset is used
        as a fallback.
        """

        if not question.file_name:
            return None

        destination = (
            Path(self.settings.download_dir)
            / question.task_id
            / question.file_name
        )

        destination.parent.mkdir(
            parents=True,
            exist_ok=True,
        )

        # Reuse an existing valid download.
        if (
            destination.exists()
            and destination.is_file()
            and destination.stat().st_size > 0
        ):
            return destination

        course_file = self._download_from_course_api(
            task_id=question.task_id,
            destination=destination,
        )

        if course_file is not None:
            return course_file

        return self._download_from_gaia_dataset(
            question=question,
            destination=destination,
        )

    def _download_from_course_api(
        self,
        task_id: str,
        destination: Path,
    ) -> Path | None:
        """Try downloading an attachment from the course API."""

        base_url = self.settings.course_api_url.rstrip("/")
        url = f"{base_url}/files/{task_id}"

        try:
            response = self.session.get(
                url=url,
                timeout=self.settings.request_timeout,
            )
        except requests.RequestException as exc:
            print(
                "Course attachment request failed "
                f"for {task_id}: {exc}"
            )
            return None

        if response.status_code == 404:
            print(
                f"Course API attachment unavailable for {task_id}; "
                "trying GAIA fallback."
            )
            return None

        try:
            response.raise_for_status()
        except requests.RequestException as exc:
            print(
                "Course attachment download failed "
                f"for {task_id}: {exc}"
            )
            return None

        content = response.content

        if not content:
            print(
                f"Course API returned an empty file for {task_id}; "
                "trying GAIA fallback."
            )
            return None

        destination.write_bytes(content)

        if destination.stat().st_size == 0:
            destination.unlink(missing_ok=True)
            return None

        return destination

    def _download_from_gaia_dataset(
        self,
        question: GaiaQuestion,
        destination: Path,
    ) -> Path:
        """Download an attachment from the official GAIA dataset."""

        token = self.settings.hf_token

        if not token:
            raise RuntimeError(
                "The course API did not provide the attachment and "
                "HF_TOKEN is missing. Add a Hugging Face read token "
                "to the .env file."
            )

        dataset_root = snapshot_download(
            repo_id=self.settings.gaia_dataset_id,
            repo_type="dataset",
            token=token,
        )

        dataset = load_dataset(
            dataset_root,
            self.settings.gaia_dataset_config,
            split=self.settings.gaia_dataset_split,
        )

        record = next(
            (
                item
                for item in dataset
                if str(item.get("task_id", "")).strip()
                == question.task_id
            ),
            None,
        )

        if record is None:
            raise FileNotFoundError(
                f"Task {question.task_id} was not found "
                "in the GAIA dataset."
            )

        source_path = self._resolve_gaia_file_path(
            dataset_root=Path(dataset_root),
            record=record,
            expected_file_name=question.file_name,
        )

        if source_path is None:
            raise FileNotFoundError(
                f"Could not locate attachment "
                f"{question.file_name!r} for task "
                f"{question.task_id} in the GAIA dataset."
            )

        destination.parent.mkdir(
            parents=True,
            exist_ok=True,
        )

        shutil.copy2(
            source_path,
            destination,
        )

        if (
            not destination.exists()
            or destination.stat().st_size == 0
        ):
            raise RuntimeError(
                f"The attachment for task {question.task_id} "
                "was copied but the destination file is empty."
            )

        return destination

    @staticmethod
    def _resolve_gaia_file_path(
        dataset_root: Path,
        record: dict[str, Any],
        expected_file_name: str,
    ) -> Path | None:
        """Resolve the attachment path inside a GAIA snapshot."""

        candidates: list[Path] = []

        raw_file_path = record.get("file_path")

        if raw_file_path:
            raw_path = Path(str(raw_file_path))

            if raw_path.is_absolute():
                candidates.append(raw_path)
            else:
                candidates.append(
                    dataset_root / raw_path
                )

        # Search by exact expected filename as a fallback.
        candidates.extend(
            dataset_root.rglob(expected_file_name)
        )

        for candidate in candidates:
            if (
                candidate.exists()
                and candidate.is_file()
                and candidate.stat().st_size > 0
            ):
                return candidate.resolve()

        return None

    def submit(
        self,
        username: str,
        agent_code: str,
        answers: list[dict[str, str]],
    ) -> dict[str, Any]:
        """Submit 20 answers to the course scoring API."""

        username = username.strip()
        agent_code = agent_code.strip()

        if not username:
            raise ValueError(
                "Hugging Face username is required."
            )

        if not agent_code:
            raise ValueError(
                "Public agent code URL is required."
            )

        if not agent_code.endswith("/tree/main"):
            raise ValueError(
                "agent_code must be a public Hugging Face Space "
                "URL ending in /tree/main."
            )

        normalized_answers = self._validate_answers(
            answers
        )

        payload = {
            "username": username,
            "agent_code": agent_code,
            "answers": normalized_answers,
        }

        base_url = self.settings.course_api_url.rstrip("/")
        url = f"{base_url}/submit"

        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "User-Agent": "Vertex-Agent/1.0",
        }

        last_error: Exception | None = None

        for attempt in range(1, 4):
            try:
                response = self.session.post(
                    url=url,
                    json=payload,
                    headers=headers,
                    timeout=(30, 300),
                    allow_redirects=False,
                )

                if response.status_code in {
                    301,
                    302,
                    303,
                    307,
                    308,
                }:
                    location = response.headers.get(
                        "Location",
                        "unknown",
                    )

                    raise RuntimeError(
                        "The scoring API redirected the submission. "
                        f"HTTP {response.status_code}. "
                        f"Location: {location}"
                    )

                if response.status_code in {401, 403}:
                    raise RuntimeError(
                        "The scoring API rejected access. "
                        f"HTTP {response.status_code}: "
                        f"{response.text[:1000]}"
                    )

                if response.status_code == 422:
                    raise RuntimeError(
                        "The scoring API rejected the payload format. "
                        f"Response: {response.text[:2000]}"
                    )

                response.raise_for_status()

                try:
                    data = response.json()
                except ValueError as exc:
                    raise RuntimeError(
                        "The scoring API returned a non-JSON response: "
                        f"{response.text[:2000]}"
                    ) from exc

                if not isinstance(data, dict):
                    raise RuntimeError(
                        "The scoring API returned an unexpected "
                        "response type."
                    )

                return data

            except requests.RequestException as exc:
                last_error = exc

                if attempt < 3:
                    delay_seconds = attempt * 3

                    print(
                        "Submission connection failed "
                        f"on attempt {attempt}/3: {exc}"
                    )
                    print(
                        f"Retrying in {delay_seconds} seconds..."
                    )

                    time.sleep(delay_seconds)
                    continue

            except RuntimeError:
                raise

        raise RuntimeError(
            "Submission failed after 3 attempts. "
            f"Last connection error: {last_error}"
        )

    @staticmethod
    def _validate_answers(
        answers: list[dict[str, str]],
    ) -> list[dict[str, str]]:
        """Validate and normalize answers before submission."""

        if not answers:
            raise ValueError(
                "Answers list cannot be empty."
            )

        if len(answers) != 20:
            raise ValueError(
                f"Expected 20 answers, but received "
                f"{len(answers)}."
            )

        normalized: list[dict[str, str]] = []
        seen_task_ids: set[str] = set()
        empty_task_ids: list[str] = []

        for index, item in enumerate(
            answers,
            start=1,
        ):
            if not isinstance(item, dict):
                raise ValueError(
                    f"Answer number {index} must be an object."
                )

            task_id = str(
                item.get("task_id", "")
            ).strip()

            submitted_answer = str(
                item.get("submitted_answer", "")
            ).strip()

            if not task_id:
                raise ValueError(
                    f"Answer number {index} is missing task_id."
                )

            if task_id in seen_task_ids:
                raise ValueError(
                    f"Duplicate task_id found: {task_id}"
                )

            seen_task_ids.add(task_id)

            if not submitted_answer:
                empty_task_ids.append(task_id)

            normalized.append(
                {
                    "task_id": task_id,
                    "submitted_answer": submitted_answer,
                }
            )

        if empty_task_ids:
            raise ValueError(
                "Submission stopped because these tasks "
                "have empty answers:\n"
                + "\n".join(empty_task_ids)
            )

        return normalized