agent.py

from __future__ import annotations

import ast
import base64
import json
import math
import operator
import os
import re
import shutil
import subprocess
import sys
import tempfile
from io import BytesIO
from pathlib import Path
from typing import Annotated, Literal
from urllib.parse import parse_qs, urlparse

import chess
import chess.engine
import cv2
import pandas as pd
import requests
import yt_dlp
from bs4 import BeautifulSoup
from ddgs import DDGS
from langchain_core.messages import (
    AIMessage,
    AnyMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from openai import OpenAI, RateLimitError
from PIL import Image as PILImage
from PIL import ImageOps
from pypdf import PdfReader
from typing_extensions import NotRequired, TypedDict
from youtube_transcript_api import YouTubeTranscriptApi


# -----------------------------------------------------------------------------
# Model configuration
# -----------------------------------------------------------------------------

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "").strip()

if not OPENAI_API_KEY:
    raise RuntimeError(
        "OPENAI_API_KEY is missing. Add it under the Hugging Face "
        "Space Settings > Variables and secrets > Secrets."
    )

TEXT_MODEL = os.getenv("TEXT_MODEL", "gpt-4.1").strip()
VISION_MODEL = os.getenv("VISION_MODEL", "gpt-4.1").strip()
AUDIO_MODEL = os.getenv("AUDIO_MODEL", "gpt-4o-mini-transcribe").strip()
OPENAI_TIMEOUT = float(os.getenv("OPENAI_TIMEOUT", "240"))
OPENAI_MAX_RETRIES = int(os.getenv("OPENAI_MAX_RETRIES", "3"))

llm = ChatOpenAI(
    model=TEXT_MODEL,
    api_key=OPENAI_API_KEY,
    temperature=0,
    max_tokens=1200,
    timeout=OPENAI_TIMEOUT,
    max_retries=OPENAI_MAX_RETRIES,
)

vision_llm = ChatOpenAI(
    model=VISION_MODEL,
    api_key=OPENAI_API_KEY,
    temperature=0,
    max_tokens=1400,
    timeout=OPENAI_TIMEOUT,
    max_retries=OPENAI_MAX_RETRIES,
)

vision_llm_chess = vision_llm

openai_client = OpenAI(
    api_key=OPENAI_API_KEY,
    timeout=OPENAI_TIMEOUT,
    max_retries=OPENAI_MAX_RETRIES,
)

print(
    "OpenAI models configured:",
    {
        "text": TEXT_MODEL,
        "vision": VISION_MODEL,
        "audio": AUDIO_MODEL,
    },
)


# -----------------------------------------------------------------------------
# General tools
# -----------------------------------------------------------------------------

@tool("web_search")
def web_search_tool(query: str) -> str:
    """Search the public web and return concise titles, URLs, and snippets."""

    query = query.strip()
    if not query:
        return "ERROR: Search query is empty."

    try:
        raw_results = list(
            DDGS().text(
                query,
                max_results=4,
            )
        )

        results = []
        for item in raw_results:
            if not isinstance(item, dict):
                continue

            title = str(item.get("title", "")).strip()
            url = str(
                item.get("href")
                or item.get("url")
                or ""
            ).strip()
            snippet = str(
                item.get("body")
                or item.get("snippet")
                or ""
            ).strip()

            if title or url or snippet:
                results.append(
                    {
                        "title": title,
                        "url": url,
                        "snippet": snippet,
                    }
                )

        if not results:
            return "ERROR: Web search returned no results."

        return json.dumps(
            results,
            ensure_ascii=False,
        )

    except Exception as error:
        return (
            "ERROR: Web search failed: "
            f"{type(error).__name__}: {error}"
        )


@tool("read_webpage")
def read_webpage(url: str) -> str:
    """Read visible text from a public webpage."""

    if not url.startswith(("http://", "https://")):
        return "ERROR: URL must begin with http:// or https://."

    try:
        response = requests.get(
            url,
            timeout=30,
            headers={
                "User-Agent": (
                    "Mozilla/5.0 (compatible; GAIAResearchAgent/1.0)"
                )
            },
        )
        response.raise_for_status()

        soup = BeautifulSoup(response.text, "html.parser")

        for element in soup(
            ["script", "style", "nav", "footer", "header", "noscript", "svg"]
        ):
            element.decompose()

        lines = [
            line.strip()
            for line in soup.get_text(separator="\n", strip=True).splitlines()
            if line.strip()
        ]

        # Remove only consecutive duplicate lines. Global de-duplication can
        # destroy repeated rows in tables.
        cleaned_lines: list[str] = []
        for line in lines:
            if not cleaned_lines or line != cleaned_lines[-1]:
                cleaned_lines.append(line)

        cleaned_text = "\n".join(cleaned_lines)
        normalized = cleaned_text.lower()

        blocked_phrases = (
            "checking your browser",
            "access denied",
            "enable javascript",
            "captcha",
        )

        if len(cleaned_text) < 100 or any(
            phrase in normalized for phrase in blocked_phrases
        ):
            return "ERROR: The webpage was blocked or contained no usable text."

        return cleaned_text[:9000]

    except requests.RequestException as error:
        return f"ERROR: Could not read webpage: {type(error).__name__}: {error}"


WIKIPEDIA_API_URL = "https://en.wikipedia.org/w/api.php"
WIKIPEDIA_HEADERS = {
    "User-Agent": "GAIA-LangGraph-Agent/1.0 (educational benchmark project)"
}


@tool("wikipedia_search")
def wikipedia_search(
    query: str,
    as_of_date: str = "2022-12-31",
) -> str:
    """
    Search English Wikipedia and return the best page's content from the
    latest revision on or before as_of_date.
    """

    try:
        search_response = requests.get(
            WIKIPEDIA_API_URL,
            params={
                "action": "query",
                "list": "search",
                "srsearch": query,
                "srlimit": 5,
                "format": "json",
                "formatversion": 2,
            },
            headers=WIKIPEDIA_HEADERS,
            timeout=30,
        )
        search_response.raise_for_status()
        results = search_response.json().get("query", {}).get("search", [])

        if not results:
            return "ERROR: No English Wikipedia page matched the query."

        query_words = set(re.findall(r"[a-z0-9]+", query.lower()))

        def score(item: dict) -> tuple[int, int]:
            title = str(item.get("title", ""))
            title_words = set(re.findall(r"[a-z0-9]+", title.lower()))
            exact = int(title.lower() == query.lower().strip())
            overlap = len(query_words & title_words)
            return exact, overlap

        page_title = max(results, key=score)["title"]

        revision_response = requests.get(
            WIKIPEDIA_API_URL,
            params={
                "action": "query",
                "prop": "revisions",
                "titles": page_title,
                "rvstart": f"{as_of_date}T23:59:59Z",
                "rvdir": "older",
                "rvlimit": 1,
                "rvprop": "ids|timestamp",
                "format": "json",
                "formatversion": 2,
            },
            headers=WIKIPEDIA_HEADERS,
            timeout=30,
        )
        revision_response.raise_for_status()
        pages = revision_response.json().get("query", {}).get("pages", [])
        revisions = pages[0].get("revisions", []) if pages else []

        if not revisions:
            return f"ERROR: No revision was found on or before {as_of_date}."

        revision_id = revisions[0]["revid"]
        revision_timestamp = revisions[0]["timestamp"]

        page_response = requests.get(
            WIKIPEDIA_API_URL,
            params={
                "action": "parse",
                "oldid": revision_id,
                "prop": "text",
                "format": "json",
                "formatversion": 2,
            },
            headers=WIKIPEDIA_HEADERS,
            timeout=30,
        )
        page_response.raise_for_status()
        html = page_response.json().get("parse", {}).get("text", "")

        if not html:
            return "ERROR: Wikipedia returned no page content."

        soup = BeautifulSoup(html, "html.parser")
        for element in soup.select(
            "script, style, sup.reference, .mw-editsection, .navbox, "
            ".vertical-navbox, .metadata"
        ):
            element.decompose()

        blocks: list[str] = []
        for element in soup.select("h2, h3, h4, p, li, tr"):
            text = " ".join(element.stripped_strings)
            if text:
                blocks.append(text)

        return json.dumps(
            {
                "title": page_title,
                "revision_timestamp": revision_timestamp,
                "content": "\n".join(blocks)[:18000],
            },
            ensure_ascii=False,
        )

    except Exception as error:
        return f"ERROR: Wikipedia lookup failed: {type(error).__name__}: {error}"


BINARY_OPERATORS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Pow: operator.pow,
    ast.Mod: operator.mod,
}
UNARY_OPERATORS = {ast.UAdd: operator.pos, ast.USub: operator.neg}


def _evaluate_math_node(node):
    if isinstance(node, ast.Expression):
        return _evaluate_math_node(node.body)
    if isinstance(node, ast.Constant):
        if not isinstance(node.value, (int, float)):
            raise ValueError("Only numbers are allowed.")
        return node.value
    if isinstance(node, ast.BinOp):
        operation_type = type(node.op)
        if operation_type not in BINARY_OPERATORS:
            raise ValueError(f"Unsupported operation: {operation_type.__name__}")
        left = _evaluate_math_node(node.left)
        right = _evaluate_math_node(node.right)
        if operation_type is ast.Pow and abs(right) > 100:
            raise ValueError("Exponent is too large.")
        return BINARY_OPERATORS[operation_type](left, right)
    if isinstance(node, ast.UnaryOp):
        operation_type = type(node.op)
        if operation_type not in UNARY_OPERATORS:
            raise ValueError("Unsupported unary operation.")
        return UNARY_OPERATORS[operation_type](_evaluate_math_node(node.operand))
    raise ValueError("Expression contains an unsupported element.")


@tool("calculator")
def calculator(expression: str) -> str:
    """Evaluate arithmetic using +, -, *, /, %, **, and parentheses."""

    if len(expression) > 200:
        return "ERROR: Calculator expression is too long."

    try:
        parsed = ast.parse(expression, mode="eval")
        return str(_evaluate_math_node(parsed))
    except Exception as error:
        return f"ERROR: Calculator failed: {type(error).__name__}: {error}"


@tool("python_executor")
def python_executor(code: str) -> str:
    """Execute short Python code for deterministic data processing."""

    if not code.strip():
        return "ERROR: No Python code was provided."
    if len(code) > 10000:
        return "ERROR: Python code is too long."

    try:
        with tempfile.TemporaryDirectory() as directory:
            completed = subprocess.run(
                [sys.executable, "-I", "-c", code],
                cwd=directory,
                capture_output=True,
                text=True,
                timeout=20,
            )

        if completed.returncode != 0:
            return f"ERROR: Python execution failed:\n{completed.stderr[:4000]}"

        output = completed.stdout.strip()
        if not output:
            return "ERROR: Python ran but printed no output."
        return output[:10000]

    except subprocess.TimeoutExpired:
        return "ERROR: Python execution exceeded 20 seconds."
    except Exception as error:
        return f"ERROR: Python execution failed: {type(error).__name__}: {error}"


# -----------------------------------------------------------------------------
# YouTube transcript tool
# -----------------------------------------------------------------------------


def extract_youtube_video_id(url: str) -> str:
    parsed_url = urlparse(url.strip())
    hostname = (parsed_url.hostname or "").lower().removeprefix("www.")
    video_id = ""

    if hostname == "youtu.be":
        video_id = parsed_url.path.strip("/").split("/")[0]
    elif hostname in {"youtube.com", "m.youtube.com", "music.youtube.com"}:
        if parsed_url.path == "/watch":
            video_id = parse_qs(parsed_url.query).get("v", [""])[0]
        elif parsed_url.path.startswith(("/shorts/", "/embed/", "/live/")):
            parts = parsed_url.path.strip("/").split("/")
            if len(parts) >= 2:
                video_id = parts[1]

    if not re.fullmatch(r"[A-Za-z0-9_-]{11}", video_id):
        raise ValueError("Could not extract a valid YouTube video ID.")
    return video_id


def format_video_timestamp(seconds: float) -> str:
    total_seconds = int(seconds)
    minutes, seconds = divmod(total_seconds, 60)
    hours, minutes = divmod(minutes, 60)
    if hours:
        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
    return f"{minutes:02d}:{seconds:02d}"


@tool("youtube_transcript")
def youtube_transcript(url: str, languages: str = "en") -> str:
    """Retrieve timestamped captions for dialogue or spoken-answer questions."""

    try:
        video_id = extract_youtube_video_id(url)
        language_codes = [x.strip() for x in languages.split(",") if x.strip()]
        transcript = YouTubeTranscriptApi().fetch(
            video_id,
            languages=language_codes or ["en"],
        )

        lines = [f"VIDEO ID: {video_id}", "TRANSCRIPT:"]
        for snippet in transcript:
            text = " ".join(snippet.text.split())
            if text:
                lines.append(f"[{format_video_timestamp(snippet.start)}] {text}")

        result = "\n".join(lines)
        return result[:18000] if result else "ERROR: No transcript was returned."

    except Exception as error:
        return f"ERROR: Transcript retrieval failed: {type(error).__name__}: {error}"


# -----------------------------------------------------------------------------
# Generic visual YouTube tool
# -----------------------------------------------------------------------------

VIDEO_MAX_FRAMES = int(os.getenv("VIDEO_MAX_FRAMES", "24"))
VIDEO_BATCH_SIZE = int(os.getenv("VIDEO_BATCH_SIZE", "8"))
VIDEO_MAX_IMAGE_SIDE = 768
VIDEO_JPEG_QUALITY = 82


def _remove_partial_video_files(output_directory: Path) -> None:
    for file_path in output_directory.glob("video.*"):
        try:
            file_path.unlink()
        except OSError:
            pass


def _find_downloaded_video(output_directory: Path) -> Path | None:
    ignored = {".part", ".ytdl", ".json", ".description"}
    files = [
        path
        for path in output_directory.glob("video.*")
        if path.is_file() and path.suffix not in ignored and path.stat().st_size > 0
    ]
    return max(files, key=lambda path: path.stat().st_size) if files else None


def download_youtube_video(url: str, output_directory: Path) -> Path:
    """Download a public YouTube video, trying several player clients."""

    output_directory.mkdir(parents=True, exist_ok=True)
    base_options = {
        "format": "best[ext=mp4][height<=480]/best[height<=480]/best",
        "outtmpl": str(output_directory / "video.%(ext)s"),
        "noplaylist": True,
        "quiet": True,
        "no_warnings": True,
        "force_ipv4": True,
        "retries": 2,
        "fragment_retries": 2,
        "socket_timeout": 30,
        "overwrites": True,
    }

    attempts = [
        ["default", "tv_simply"],
        ["web_safari", "tv_simply"],
    ]
    errors: list[str] = []

    for clients in attempts:
        _remove_partial_video_files(output_directory)
        options = dict(base_options)
        options["extractor_args"] = {"youtube": {"player_client": clients}}

        try:
            with yt_dlp.YoutubeDL(options) as downloader:
                downloader.download([url])
            downloaded = _find_downloaded_video(output_directory)
            if downloaded:
                return downloaded
        except Exception as error:
            errors.append(f"{clients}: {type(error).__name__}: {error}")

    raise RuntimeError("All YouTube download attempts failed: " + " | ".join(errors))


def resize_video_frame(frame, maximum_side: int = VIDEO_MAX_IMAGE_SIDE):
    height, width = frame.shape[:2]
    longest = max(width, height)
    if longest <= maximum_side:
        return frame
    scale = maximum_side / longest
    return cv2.resize(
        frame,
        (max(1, int(width * scale)), max(1, int(height * scale))),
        interpolation=cv2.INTER_AREA,
    )


def sample_video_frames(
    video_path: Path,
    maximum_frames: int = VIDEO_MAX_FRAMES,
) -> list[dict]:
    capture = cv2.VideoCapture(str(video_path))

    try:
        if not capture.isOpened():
            raise ValueError("OpenCV could not open the video.")

        fps = float(capture.get(cv2.CAP_PROP_FPS))
        frame_count = float(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        if fps <= 0 or frame_count <= 0:
            raise ValueError("Could not determine video duration.")

        duration = frame_count / fps
        sample_count = min(maximum_frames, max(12, math.ceil(duration)))
        final_timestamp = max(duration - 0.05, 0.0)
        timestamps = [
            index * final_timestamp / max(sample_count - 1, 1)
            for index in range(sample_count)
        ]

        sampled: list[dict] = []
        for timestamp in timestamps:
            capture.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
            success, frame = capture.read()
            if not success:
                continue
            frame = resize_video_frame(frame)
            encoded_success, encoded = cv2.imencode(
                ".jpg",
                frame,
                [int(cv2.IMWRITE_JPEG_QUALITY), VIDEO_JPEG_QUALITY],
            )
            if not encoded_success:
                continue
            sampled.append(
                {
                    "timestamp_seconds": round(timestamp, 3),
                    "image_base64": base64.b64encode(encoded.tobytes()).decode(),
                }
            )
        return sampled
    finally:
        capture.release()


def _model_content_to_text(content) -> str:
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        return "\n".join(
            str(block.get("text", ""))
            for block in content
            if isinstance(block, dict) and block.get("text")
        ).strip()
    return str(content).strip()


def _extract_json_object(text: str) -> dict:
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end <= start:
        raise ValueError("The model did not return a JSON object.")
    return json.loads(text[start : end + 1])


def _safe_float(value, default=None):
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def analyze_video_frame_batch(
    frame_batch: list[dict],
    question: str,
) -> list[dict]:
    prompt = f"""
Analyze each labeled frame independently for the original visual question.

ORIGINAL QUESTION:
{question}

For every frame:
1. Decide whether it contains relevant visible evidence.
2. Describe only what is visibly present.
3. Never combine counts or objects across timestamps.
4. For a count question, put the value supported by that frame in numeric_value.
5. For an identification, color, text, object, person, animal, action, place,
   or event question, put the possible answer in candidate_answer.
6. Use null when the frame does not support a value.
7. Be conservative when evidence is unclear.

Return JSON only:
{{
  "frames": [
    {{
      "frame_label": "FRAME 1",
      "relevant": true,
      "observation": "visible evidence",
      "candidate_answer": null,
      "numeric_value": null,
      "confidence": 0.0
    }}
  ]
}}
""".strip()

    content: list[dict] = [{"type": "text", "text": prompt}]
    for index, frame in enumerate(frame_batch, start=1):
        content.append(
            {
                "type": "text",
                "text": f"FRAME {index} — {frame['timestamp_seconds']:.2f} seconds",
            }
        )
        content.append(
            {
                "type": "image_url",
                "image_url": {
                    "url": "data:image/jpeg;base64," + frame["image_base64"]
                },
            }
        )

    response = vision_llm.invoke([HumanMessage(content=content)])
    parsed = _extract_json_object(_model_content_to_text(response.content))
    returned = {
        item.get("frame_label"): item
        for item in parsed.get("frames", [])
        if isinstance(item, dict)
    }

    observations: list[dict] = []
    for index, frame in enumerate(frame_batch, start=1):
        result = returned.get(f"FRAME {index}", {})
        observations.append(
            {
                "timestamp_seconds": frame["timestamp_seconds"],
                "relevant": bool(result.get("relevant", False)),
                "observation": str(result.get("observation", "")).strip(),
                "candidate_answer": result.get("candidate_answer"),
                "numeric_value": _safe_float(result.get("numeric_value")),
                "confidence": _safe_float(result.get("confidence"), 0.0),
            }
        )
    return observations


def synthesize_video_answer(question: str, observations: list[dict]) -> dict:
    relevant = [item for item in observations if item.get("relevant")]
    if not relevant:
        return {"answer": "Unknown", "evidence_timestamps": [], "confidence": 0.0}

    prompt = f"""
Answer the original question using only these timestamped visual observations.

ORIGINAL QUESTION:
{question}

OBSERVATIONS:
{json.dumps(relevant[:60], ensure_ascii=False)}

Rules:
- For highest/maximum/most simultaneously, use the largest value from one timestamp.
- For lowest/minimum, use the smallest value from one timestamp.
- For first, use the earliest relevant timestamp.
- For last, use the latest relevant timestamp.
- Do not add values across timestamps.
- Return Unknown when evidence is insufficient.

Return JSON only:
{{"answer": "concise answer", "evidence_timestamps": [0.0], "confidence": 0.0}}
""".strip()

    response = vision_llm.invoke([HumanMessage(content=prompt)])
    result = _extract_json_object(_model_content_to_text(response.content))
    return {
        "answer": str(result.get("answer", "Unknown")).strip(),
        "evidence_timestamps": result.get("evidence_timestamps", []),
        "confidence": _safe_float(result.get("confidence"), 0.0),
    }


@tool("youtube_visual_analysis")
def youtube_visual_analysis(url: str, question: str) -> str:
    """Analyze objects, counts, text, colors, actions, and events visible in video."""

    try:
        with tempfile.TemporaryDirectory() as directory:
            video_path = download_youtube_video(url, Path(directory))
            sampled_frames = sample_video_frames(video_path, VIDEO_MAX_FRAMES)
            if not sampled_frames:
                return json.dumps({"error": "No video frames could be extracted."})

            observations: list[dict] = []
            batch_errors: list[str] = []
            completed_batches = 0

            for batch_start in range(0, len(sampled_frames), VIDEO_BATCH_SIZE):
                batch_number = batch_start // VIDEO_BATCH_SIZE + 1
                batch = sampled_frames[batch_start : batch_start + VIDEO_BATCH_SIZE]
                try:
                    observations.extend(analyze_video_frame_batch(batch, question))
                    completed_batches += 1
                except RateLimitError as error:
                    return json.dumps(
                        {
                            "error": "Vision-model API rate limit reached.",
                            "stage": f"frame-analysis batch {batch_number}",
                            "provider_message": str(error),
                        },
                        ensure_ascii=False,
                    )
                except Exception as error:
                    batch_errors.append(
                        f"Batch {batch_number}: {type(error).__name__}: {error}"
                    )

            if not observations:
                return json.dumps(
                    {
                        "error": "No video frames were successfully analyzed.",
                        "batch_errors": batch_errors,
                    },
                    ensure_ascii=False,
                )

            normalized = question.lower()
            numeric = [
                item
                for item in observations
                if item.get("relevant") and item.get("numeric_value") is not None
            ]
            final_result = None

            if numeric and any(
                phrase in normalized
                for phrase in (
                    "highest number",
                    "maximum number",
                    "largest number",
                    "most simultaneously",
                )
            ):
                best = max(numeric, key=lambda item: item["numeric_value"])
                value = best["numeric_value"]
                value = int(value) if float(value).is_integer() else value
                final_result = {
                    "answer": str(value),
                    "evidence_timestamps": [best["timestamp_seconds"]],
                    "confidence": best.get("confidence", 0.0),
                }
            elif numeric and any(
                phrase in normalized
                for phrase in ("lowest number", "minimum number", "smallest number")
            ):
                best = min(numeric, key=lambda item: item["numeric_value"])
                value = best["numeric_value"]
                value = int(value) if float(value).is_integer() else value
                final_result = {
                    "answer": str(value),
                    "evidence_timestamps": [best["timestamp_seconds"]],
                    "confidence": best.get("confidence", 0.0),
                }

            if final_result is None:
                final_result = synthesize_video_answer(question, observations)

            return json.dumps(
                {
                    **final_result,
                    "frames_analyzed": len(observations),
                    "frames_sampled": len(sampled_frames),
                    "batches_completed": completed_batches,
                    "batch_errors": batch_errors,
                },
                ensure_ascii=False,
            )

    except RateLimitError as error:
        return json.dumps(
            {
                "error": "Vision-model API rate limit reached.",
                "provider_message": str(error),
            }
        )
    except Exception as error:
        return json.dumps(
            {
                "error": (
                    "YouTube visual analysis failed: "
                    f"{type(error).__name__}: {error}"
                )
            },
            ensure_ascii=False,
        )


# -----------------------------------------------------------------------------
# Agent state and attachment routing
# -----------------------------------------------------------------------------

RouteType = Literal[
    "reasoning",
    "audio",
    "image",
    "chess",
    "spreadsheet",
    "python_file",
    "pdf",
]


class AgentState(TypedDict):
    question: str
    messages: Annotated[list[AnyMessage], add_messages]
    route: NotRequired[RouteType]
    input_file: NotRequired[str]
    attachment_content: NotRequired[str]
    final_answer: NotRequired[str]
    error: NotRequired[str]


AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".flac", ".ogg"}
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
SPREADSHEET_EXTENSIONS = {".csv", ".xlsx", ".xls", ".xlsm"}


def router_node(state: AgentState) -> dict:
    question = state["question"].lower()
    input_file = state.get("input_file")

    if not input_file:
        return {"route": "reasoning"}

    extension = Path(input_file).suffix.lower()
    if extension in AUDIO_EXTENSIONS:
        return {"route": "audio"}
    if extension in IMAGE_EXTENSIONS:
        chess_keywords = (
            "chess",
            "black's turn",
            "white's turn",
            "algebraic notation",
            "checkmate",
        )
        return {
            "route": "chess" if any(x in question for x in chess_keywords) else "image"
        }
    if extension in SPREADSHEET_EXTENSIONS:
        return {"route": "spreadsheet"}
    if extension == ".py":
        return {"route": "python_file"}
    if extension == ".pdf":
        return {"route": "pdf"}
    return {"route": "reasoning"}


def choose_route(state: AgentState) -> RouteType:
    return state.get("route", "reasoning")


# -----------------------------------------------------------------------------
# Image helpers and nodes
# -----------------------------------------------------------------------------

MAX_IMAGE_SIDE = 768
JPEG_QUALITY = 85


def prepare_image_for_vlm(file_path: Path) -> tuple[str, str, tuple[int, int]]:
    with PILImage.open(file_path) as image:
        image = ImageOps.exif_transpose(image)

        if image.mode in ("RGBA", "LA"):
            background = PILImage.new("RGB", image.size, "white")
            background.paste(image.convert("RGB"), mask=image.getchannel("A"))
            image = background
        elif image.mode == "P" and "transparency" in image.info:
            image = image.convert("RGBA")
            background = PILImage.new("RGB", image.size, "white")
            background.paste(image.convert("RGB"), mask=image.getchannel("A"))
            image = background
        else:
            image = image.convert("RGB")

        image.thumbnail((MAX_IMAGE_SIDE, MAX_IMAGE_SIDE), PILImage.Resampling.LANCZOS)
        resized_size = image.size
        buffer = BytesIO()
        image.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True)
        image_base64 = base64.b64encode(buffer.getvalue()).decode()

    return image_base64, "image/jpeg", resized_size


def image_node(state: AgentState) -> dict:
    input_file = state.get("input_file")
    if not input_file:
        return {"error": "No image file was supplied."}

    file_path = Path(input_file)
    try:
        image_base64, mime_type, resized_size = prepare_image_for_vlm(file_path)
        question = state.get("question", "Describe the image.").strip()
        prompt = f"""
Analyze the attached image for the original question.

ORIGINAL QUESTION:
{question}

Extract only relevant visible evidence, including readable text, numbers,
symbols, labels, objects, positions, tables, and chart values. Do not invent
unclear details and do not use outside knowledge.
""".strip()
        response = vision_llm.invoke(
            [
                HumanMessage(
                    content=[
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{image_base64}"
                            },
                        },
                    ]
                )
            ]
        )
        analysis = _model_content_to_text(response.content)
        if not analysis:
            raise ValueError("The vision model returned no image analysis.")
        return {
            "attachment_content": (
                "IMAGE ANALYSIS\n\n"
                f"FILE NAME: {file_path.name}\n"
                f"RESIZED DIMENSIONS: {resized_size[0]} x {resized_size[1]}\n\n"
                f"VISUAL CONTENT:\n{analysis}"
            )
        }
    except Exception as error:
        message = f"Image analysis failed: {type(error).__name__}: {error}"
        return {"attachment_content": message, "error": message}


# -----------------------------------------------------------------------------
# Audio node
# -----------------------------------------------------------------------------


def audio_node(state: AgentState) -> dict:
    input_file = state.get("input_file")
    if not input_file:
        return {"error": "No audio file was supplied."}

    file_path = Path(input_file)

    try:
        with file_path.open("rb") as audio_file:
            transcription = openai_client.audio.transcriptions.create(
                model=AUDIO_MODEL,
                file=audio_file,
                language="en",
                prompt=(
                    "Transcribe accurately. Preserve names, numbers, page "
                    "numbers, ingredient names, and punctuation."
                ),
                response_format="text",
            )

        if isinstance(transcription, str):
            transcript = transcription.strip()
        else:
            transcript = str(
                getattr(transcription, "text", "")
            ).strip()

        if not transcript:
            raise ValueError("The transcription API returned no text.")

        return {
            "attachment_content": (
                "AUDIO TRANSCRIPTION\n\n"
                f"FILE NAME: {file_path.name}\n"
                f"TRANSCRIPTION MODEL: {AUDIO_MODEL}\n\n"
                f"TRANSCRIPT:\n{transcript}"
            )
        }

    except RateLimitError as error:
        message = (
            "Audio transcription failed because the OpenAI rate limit "
            f"was reached: {error}"
        )
        return {"attachment_content": message, "error": message}

    except Exception as error:
        message = (
            "Audio transcription failed: "
            f"{type(error).__name__}: {error}"
        )
        return {"attachment_content": message, "error": message}


# -----------------------------------------------------------------------------
# Chess node
# -----------------------------------------------------------------------------

def find_stockfish() -> str | None:
    candidates = [
        os.getenv("STOCKFISH_PATH", "").strip(),
        shutil.which("stockfish"),
        "/usr/games/stockfish",
        "/usr/local/bin/stockfish",
        "/usr/bin/stockfish",
    ]

    for candidate in candidates:
        if candidate and Path(candidate).is_file():
            return candidate

    return None


STOCKFISH_PATH = find_stockfish()


def detect_side_from_question(question: str):
    normalized = question.lower().replace("’", "'")
    if any(x in normalized for x in ("black to move", "black's turn", "move for black")):
        return chess.BLACK
    if any(x in normalized for x in ("white to move", "white's turn", "move for white")):
        return chess.WHITE
    return None


def chess_node(state: AgentState) -> dict:
    input_file = state.get("input_file")
    if not input_file:
        return {"error": "No chess image was supplied."}
    if not STOCKFISH_PATH:
        message = "Stockfish is not installed."
        return {"attachment_content": message, "error": message}

    file_path = Path(input_file)
    question = state.get("question", "").strip()

    try:
        image_base64, mime_type, resized_size = prepare_image_for_vlm(file_path)
        explicit_turn = detect_side_from_question(question)
        turn_instruction = (
            "Set side_to_move to black."
            if explicit_turn == chess.BLACK
            else "Set side_to_move to white."
            if explicit_turn == chess.WHITE
            else "Determine the side to move from the image."
        )

        prompt = f"""
Reconstruct the attached chessboard exactly.

ORIGINAL QUESTION:
{question}

Inspect all 64 squares and board labels. Do not calculate a move.
{turn_instruction}

Return JSON only:
{{
  "white_pieces": ["Kg1"],
  "black_pieces": ["Kg8"],
  "fen": "complete FEN",
  "side_to_move": "black or white",
  "orientation": "black or white",
  "confidence": 0.0
}}
""".strip()

        response = vision_llm_chess.invoke(
            [
                HumanMessage(
                    content=[
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{image_base64}"
                            },
                        },
                    ]
                )
            ]
        )
        result = _extract_json_object(_model_content_to_text(response.content))
        fen = str(result.get("fen", "")).strip()
        if not fen:
            raise ValueError("The vision model did not return a FEN.")

        board = chess.Board(fen)
        if explicit_turn is not None:
            board.turn = explicit_turn
            fen = board.fen()
        if not board.is_valid() or board.is_game_over():
            raise ValueError(f"Invalid or finished reconstructed position: {fen}")

        engine = chess.engine.SimpleEngine.popen_uci(STOCKFISH_PATH, timeout=30.0)
        try:
            engine_result = engine.play(board, chess.engine.Limit(depth=18))
            if engine_result.move is None:
                raise ValueError("Stockfish did not return a move.")
            san = board.san(engine_result.move)
            uci = engine_result.move.uci()
        finally:
            engine.quit()

        return {
            "attachment_content": (
                "CHESS POSITION ANALYSIS\n\n"
                f"FILE NAME: {file_path.name}\n"
                f"IMAGE DIMENSIONS: {resized_size[0]} x {resized_size[1]}\n"
                f"FEN: {fen}\n"
                f"BEST MOVE IN SAN: {san}\n"
                f"BEST MOVE IN UCI: {uci}\n"
            )
        }
    except Exception as error:
        message = f"Chess processing failed: {type(error).__name__}: {error}"
        return {"attachment_content": message, "error": message}


# -----------------------------------------------------------------------------
# Spreadsheet, Python, and PDF nodes
# -----------------------------------------------------------------------------


def _clean_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame:
    cleaned = dataframe.copy().replace(r"^\s*$", pd.NA, regex=True)
    return cleaned.dropna(axis=0, how="all").dropna(axis=1, how="all")


def _dataframe_to_text(sheet_name: str, dataframe: pd.DataFrame) -> str:
    dataframe = _clean_dataframe(dataframe)
    rows, columns = dataframe.shape
    section = [
        f"SHEET NAME: {sheet_name}",
        f"ROWS: {rows}",
        f"COLUMNS: {columns}",
        "COLUMN NAMES: " + " | ".join(str(x) for x in dataframe.columns),
        "SHEET DATA:",
        dataframe.to_csv(index=False, na_rep=""),
    ]

    numeric = dataframe.apply(pd.to_numeric, errors="coerce")
    totals = numeric.sum(min_count=1).dropna()
    if not totals.empty:
        section.append("NUMERIC COLUMN TOTALS:")
        for column, total in totals.items():
            section.append(f"{column}: {total}")

    return "\n".join(section)


def spreadsheet_node(state: AgentState) -> dict:
    input_file = state.get("input_file")
    if not input_file:
        return {"error": "No spreadsheet file was supplied."}

    file_path = Path(input_file)
    try:
        if file_path.suffix.lower() == ".csv":
            sheets = {"CSV": pd.read_csv(file_path, dtype=object, keep_default_na=False)}
        else:
            sheets = pd.read_excel(
                file_path,
                sheet_name=None,
                dtype=object,
                keep_default_na=False,
            )

        content = "\n\n".join(
            _dataframe_to_text(name, frame) for name, frame in sheets.items()
        )
        return {
            "attachment_content": (
                "SPREADSHEET INFORMATION\n\n"
                f"FILE NAME: {file_path.name}\n\n{content[:30000]}"
            )
        }
    except Exception as error:
        message = f"Spreadsheet processing failed: {type(error).__name__}: {error}"
        return {"attachment_content": message, "error": message}


def python_file_node(state: AgentState) -> dict:
    input_file = state.get("input_file")
    if not input_file:
        return {"error": "No Python file was supplied."}

    file_path = Path(input_file)
    try:
        try:
            source = file_path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            source = file_path.read_text(encoding="latin-1")

        tree = ast.parse(source)
        functions = sorted(
            {
                node.name
                for node in ast.walk(tree)
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
            }
        )
        classes = sorted(
            {node.name for node in ast.walk(tree) if isinstance(node, ast.ClassDef)}
        )

        return {
            "attachment_content": (
                "PYTHON FILE INFORMATION\n\n"
                f"FILE NAME: {file_path.name}\n"
                f"FUNCTIONS: {', '.join(functions) or 'None'}\n"
                f"CLASSES: {', '.join(classes) or 'None'}\n\n"
                "ANALYSIS INSTRUCTION: Trace execution from __main__ to the final "
                "printed value. Follow loops, recursion, exceptions, generators, "
                "returns, mutations, and stopping conditions.\n\n"
                f"SOURCE CODE:\n{source[:30000]}"
            )
        }
    except Exception as error:
        message = f"Python file processing failed: {type(error).__name__}: {error}"
        return {"attachment_content": message, "error": message}


def pdf_node(state: AgentState) -> dict:
    input_file = state.get("input_file")
    if not input_file:
        return {"error": "No PDF file was supplied."}

    file_path = Path(input_file)
    try:
        reader = PdfReader(str(file_path))
        pages: list[str] = []
        for page_number, page in enumerate(reader.pages[:50], start=1):
            try:
                text = page.extract_text(extraction_mode="layout") or ""
            except TypeError:
                text = page.extract_text() or ""
            pages.append(f"--- PAGE {page_number} ---\n{text.strip()}")

        full_text = "\n\n".join(pages)
        if not full_text.strip():
            raise ValueError("No text could be extracted from the PDF.")
        return {
            "attachment_content": (
                "PDF INFORMATION\n\n"
                f"FILE NAME: {file_path.name}\n"
                f"TOTAL PAGES: {len(reader.pages)}\n\n"
                f"PDF CONTENT:\n{full_text[:15000]}"
            )
        }
    except Exception as error:
        message = f"PDF processing failed: {type(error).__name__}: {error}"
        return {"attachment_content": message, "error": message}


# -----------------------------------------------------------------------------
# Reasoning and graph
# -----------------------------------------------------------------------------

GENERAL_TOOLS = [
    web_search_tool,
    wikipedia_search,
    read_webpage,
    calculator,
    python_executor,
    youtube_transcript,
    youtube_visual_analysis,
]

WIKIPEDIA_TOOLS = [wikipedia_search, calculator, python_executor]

llm_with_tools = llm.bind_tools(GENERAL_TOOLS)
wikipedia_llm_with_tools = llm.bind_tools(WIKIPEDIA_TOOLS)
tool_node = ToolNode(GENERAL_TOOLS, handle_tool_errors=True)

MAX_TOOL_RESULTS = 6

GAIA_REASONING_PROMPT = """
You are solving a GAIA benchmark question. Use direct reasoning and the
minimum necessary tool calls.

1. Return exactly one block: <answer>YOUR ANSWER</answer>. Put no text outside it.
2. Follow exact formatting: number, name, IOC code, comma-separated list,
   alphabetical order, decimals, capitalization, punctuation, or chess SAN.
3. Solve reversed text, wordplay, simple logic, and short transformations
   directly without tools.
4. Treat attachment content as the primary source and preserve exact values.
5. For web research, use focused web_search and open a relevant result with
   read_webpage. Do not answer from snippets or blocked pages.
6. When the question mentions Wikipedia, use wikipedia_search with the main
   topic. For a latest-2022 request use as_of_date=2022-12-31. Do not switch to
   general web search unless the Wikipedia tool returns an error.
7. Use youtube_transcript for speech, dialogue, quotations, and what someone
   said. Use youtube_visual_analysis for visible objects, animals, people,
   colors, text, actions, counts, timestamps, and simultaneous events.
8. After a successful YouTube tool result, use its evidence instead of a web
   guess. Never add counts from different timestamps for a simultaneous count.
9. Use calculator for arithmetic and python_executor for sorting, filtering,
   counting, tables, comparisons, and multi-step verification.
10. For attached Python code, trace the actual entry point and final printed
    output through recursion, loops, exceptions, generators, and returns.
11. When attachment content contains BEST MOVE IN SAN, copy it exactly.
12. For counting and list questions, identify every qualifying record, apply
    every condition, verify dates/categories, then count or sort.
13. Never invent an answer because a tool failed. Avoid repeating the same
    failing call. Once evidence is sufficient, stop using tools.
14. Before answering, verify exact question, conditions, ordering, spelling,
    capitalization, symbols, units, and decimal places.
""".strip()


def reasoning_node(state: AgentState) -> dict:
    question = state["question"].strip()
    attachment_content = state.get("attachment_content", "").strip()
    messages = list(state.get("messages", []))
    if not messages:
        messages = [HumanMessage(content=question)]

    tool_result_count = sum(isinstance(message, ToolMessage) for message in messages)
    system_content = f"{GAIA_REASONING_PROMPT}\n\nORIGINAL QUESTION:\n{question}"
    if attachment_content:
        system_content += (
            "\n\nCONTENT EXTRACTED FROM THE ATTACHMENT:\n" + attachment_content
        )

    if tool_result_count >= MAX_TOOL_RESULTS:
        system_content += (
            "\n\nThe tool-use budget is exhausted. Do not call another tool. "
            "Use the reliable evidence already available and return the answer now."
        )
        selected_model = llm
    elif "wikipedia" in question.lower():
        selected_model = wikipedia_llm_with_tools
    else:
        selected_model = llm_with_tools

    response = selected_model.invoke(
        [SystemMessage(content=system_content), *messages]
    )
    return {"messages": [response]}


def content_to_text(content) -> str:
    return _model_content_to_text(content)


def final_answer_formatter(state: AgentState) -> dict:
    for message in reversed(state.get("messages", [])):
        if not isinstance(message, AIMessage) or getattr(message, "tool_calls", None):
            continue
        text = content_to_text(message.content)
        if not text:
            continue
        tagged = re.search(
            r"<answer>\s*(.*?)\s*</answer>",
            text,
            flags=re.IGNORECASE | re.DOTALL,
        )
        answer = tagged.group(1).strip() if tagged else text
        answer = re.sub(
            r"^(final\s+answer|answer|result)\s*:\s*",
            "",
            answer,
            flags=re.IGNORECASE,
        )
        answer = re.sub(r"</?answer>", "", answer, flags=re.IGNORECASE)
        return {"final_answer": answer.strip("` \n")}

    return {"final_answer": "", "error": "No completed AI answer was found."}


def route_after_processor(state: AgentState) -> Literal["reason", "stop"]:
    return "stop" if state.get("error") else "reason"


def route_after_reasoning(
    state: AgentState,
) -> Literal["use_tools", "format_answer"]:
    last_message = state["messages"][-1]
    return "use_tools" if getattr(last_message, "tool_calls", None) else "format_answer"


graph_builder = StateGraph(AgentState)
graph_builder.add_node("router", router_node)
graph_builder.add_node("audio_node", audio_node)
graph_builder.add_node("image_node", image_node)
graph_builder.add_node("chess_node", chess_node)
graph_builder.add_node("spreadsheet_node", spreadsheet_node)
graph_builder.add_node("python_file_node", python_file_node)
graph_builder.add_node("pdf_node", pdf_node)
graph_builder.add_node("reasoning_node", reasoning_node)
graph_builder.add_node("tools", tool_node)
graph_builder.add_node("final_answer_formatter", final_answer_formatter)

graph_builder.add_edge(START, "router")
graph_builder.add_conditional_edges(
    "router",
    choose_route,
    {
        "reasoning": "reasoning_node",
        "audio": "audio_node",
        "image": "image_node",
        "chess": "chess_node",
        "spreadsheet": "spreadsheet_node",
        "python_file": "python_file_node",
        "pdf": "pdf_node",
    },
)

for processor in (
    "audio_node",
    "image_node",
    "chess_node",
    "spreadsheet_node",
    "python_file_node",
    "pdf_node",
):
    graph_builder.add_conditional_edges(
        processor,
        route_after_processor,
        {"reason": "reasoning_node", "stop": END},
    )

graph_builder.add_conditional_edges(
    "reasoning_node",
    route_after_reasoning,
    {"use_tools": "tools", "format_answer": "final_answer_formatter"},
)
graph_builder.add_edge("tools", "reasoning_node")
graph_builder.add_edge("final_answer_formatter", END)

gaia_graph = graph_builder.compile()


def clean_answer(answer: str) -> str:
    answer = str(answer or "").strip()
    match = re.search(
        r"<answer>\s*(.*?)\s*</answer>",
        answer,
        flags=re.IGNORECASE | re.DOTALL,
    )
    if match:
        answer = match.group(1).strip()
    return re.sub(
        r"^(final\s+answer|answer|result)\s*:\s*",
        "",
        answer,
        flags=re.IGNORECASE,
    ).strip()


class GaiaAgent:
    """Wrapper called by app.py and the local dry-run script."""

    def __init__(self):
        self.graph = gaia_graph

    def health_check(self) -> dict:
        """Make one small paid request to validate the configured text model."""

        response = llm.invoke(
            [HumanMessage(content="Return exactly the word OK and nothing else.")]
        )
        text = _model_content_to_text(response.content).strip()

        return {
            "text_model": TEXT_MODEL,
            "vision_model": VISION_MODEL,
            "audio_model": AUDIO_MODEL,
            "text_response": text,
            "stockfish_available": bool(STOCKFISH_PATH),
            "ffmpeg_available": bool(shutil.which("ffmpeg")),
        }

    def __call__(
        self,
        question: str,
        input_file: str | None = None,
    ) -> str:
        state: AgentState = {"question": question, "messages": []}
        if input_file:
            state["input_file"] = input_file

        result = self.graph.invoke(state, config={"recursion_limit": 25})
        answer = clean_answer(result.get("final_answer", ""))
        if not answer:
            print("Agent returned no answer. Error:", result.get("error", ""))
        return answer