"""GAIA agent built with LangGraph's ReAct loop over a HF Inference model.

The agent is given a small toolset (web search, Wikipedia, page reader, task
file reader, calculator) and a system prompt that enforces the GAIA exact-match
answer format. Call `GaiaAgent()(question, task_id)` to get a clean answer
string ready for submission.
"""

from __future__ import annotations

import os
import re
import time

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langgraph.prebuilt import create_react_agent

from tools import TOOLS

# Lower-cased tool names, used to detect a model that echoes a tool name as its
# final answer instead of producing a real answer.
_TOOL_NAMES = {t.name.lower() for t in TOOLS}

# Backend selection:
#   "hf"     -> Hugging Face Inference API (needs credits)
#   "groq"   -> Groq API (free tier, fast, strong tool-calling; needs GROQ_API_KEY)
#   "ollama" -> local model via Ollama (free, offline, no quota)
GAIA_BACKEND = os.getenv("GAIA_BACKEND", "hf").lower()

# Default model on HF Inference (OpenAI-style tool calling). Override with
# GAIA_MODEL_ID.
DEFAULT_MODEL_ID = os.getenv("GAIA_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")

# Default Groq model (tool-capable). Override with GAIA_GROQ_MODEL.
DEFAULT_GROQ_MODEL = os.getenv("GAIA_GROQ_MODEL", "openai/gpt-oss-120b")

# Local Ollama model used when GAIA_BACKEND=ollama. Must be a tool-capable
# model (e.g. qwen2:7b, qwen2.5:7b, llama3.1:8b). Override with GAIA_OLLAMA_MODEL.
DEFAULT_OLLAMA_MODEL = os.getenv("GAIA_OLLAMA_MODEL", "qwen2:7b")

# GAIA-style answer formatting rules, minus the literal "FINAL ANSWER" text
# (the course asks submissions to omit that phrase and reply with the answer
# only).
SYSTEM_PROMPT = """You are a precise general AI assistant solving questions from \
the GAIA benchmark.

Work step by step. Use the available tools to look up facts, read web pages, \
read attached task files, and do exact arithmetic. Do not guess when a tool can \
get you the answer.

When you have solved the question, reply with ONLY the final answer and nothing \
else. Do not add explanation, restating, punctuation, or any prefix such as \
"Answer:".

Follow these formatting rules for the answer:
- If the answer is a number, write digits only. Do not use thousands separators \
and do not include units or symbols (no $, %, etc.) unless the question asks for \
them.
- If the answer is a string, do not use articles or abbreviations, and write \
digits in plain text unless told otherwise. Use the minimum number of words.
- If the answer is a comma separated list, apply the rules above to each \
element and separate elements with ", ".

Extra care:
- Expand abbreviations fully when the question asks for an unabbreviated form \
(e.g. write "Saint Petersburg", not "St. Petersburg").
- If the question text looks reversed, scrambled, or encoded, decode it first, \
then answer what it actually asks (and match the form it requests).
- If a referenced attachment cannot be downloaded, still give your single best \
concrete answer in the required format. Never reply with an apology or a request \
for the file.

The answer is graded by EXACT string match, so be concise and exact."""


# Phrases / prefixes the model sometimes prepends despite instructions; we strip
# them so the submitted answer is clean.
_PREFIX_RE = re.compile(
    r"^\s*(final answer|answer|the answer is)\s*[:\-]?\s*", re.IGNORECASE
)


def _build_llm(temperature: float = 0.1):
    """Construct the chat model for the selected backend.

    GAIA_BACKEND=ollama -> local model via Ollama (free, no credits).
    GAIA_BACKEND=groq   -> Groq API (free tier).
    GAIA_BACKEND=hf (default) -> Hugging Face Inference API.
    """
    if GAIA_BACKEND == "ollama":
        from langchain_ollama import ChatOllama

        return ChatOllama(
            model=DEFAULT_OLLAMA_MODEL,
            temperature=temperature,
            num_predict=1024,
        )

    if GAIA_BACKEND == "groq":
        from langchain_groq import ChatGroq

        return ChatGroq(
            model=DEFAULT_GROQ_MODEL,
            temperature=temperature,
            max_tokens=1024,
            api_key=os.getenv("GROQ_API_KEY"),
        )

    token = (
        os.getenv("HF_TOKEN")
        or os.getenv("HUGGINGFACEHUB_API_TOKEN")
        or os.getenv("HUGGINGFACE_TOKEN")
    )
    endpoint = HuggingFaceEndpoint(
        repo_id=DEFAULT_MODEL_ID,
        task="text-generation",
        huggingfacehub_api_token=token,
        temperature=temperature,
        max_new_tokens=1024,
        timeout=120,
    )
    return ChatHuggingFace(llm=endpoint)


def _clean_answer(text: str) -> str:
    """Normalize the model's final message into a submittable answer string."""
    text = (text or "").strip()
    # Strip surrounding code fences / quotes the model sometimes adds.
    text = text.strip("`").strip().strip('"').strip("'").strip()
    # Drop a leading "Final answer:" style prefix if present.
    text = _PREFIX_RE.sub("", text)
    # Collapse internal whitespace/newlines.
    text = re.sub(r"\s+", " ", text).strip()
    return text


# Marker of Groq/OpenAI-style server-side tool-call validation failures, which
# we recover from by retrying at a higher temperature.
_TOOL_FAIL_MARKERS = ("tool_use_failed", "tool call validation", "failed to call a function")

# Markers of per-minute rate limiting (e.g. Groq free tier TPM). We wait out the
# window and retry rather than failing the question.
_RATE_LIMIT_MARKERS = ("rate_limit", "tokens per minute", "request too large", "429", "413")
_RATE_LIMIT_WAIT = int(os.getenv("GAIA_RATE_LIMIT_WAIT", "30"))  # seconds


class GaiaAgent:
    """Callable agent that answers a single GAIA question."""

    def __init__(self) -> None:
        # Primary graph is near-deterministic; the retry graph samples more so a
        # malformed tool call (which some providers hard-reject) can be retried.
        self.llm = _build_llm(temperature=0.1)
        self.graph = create_react_agent(self.llm, TOOLS)
        self._retry_llm = None
        self._retry_graph = None

    def _retry(self):
        """Lazily build a higher-temperature graph used for retries."""
        if self._retry_graph is None:
            self._retry_llm = _build_llm(temperature=0.6)
            self._retry_graph = create_react_agent(self._retry_llm, TOOLS)
        return self._retry_graph

    def __call__(
        self,
        question: str,
        task_id: str | None = None,
        file_name: str | None = None,
    ) -> str:
        prompt = question
        # Only point at the file tool when the task actually has an attachment,
        # otherwise the model wastes a (often malformed) tool call on every task.
        if task_id and file_name:
            prompt = (
                f"{question}\n\n"
                f"(This task has an attached file. Call read_task_file with "
                f"task_id='{task_id}' to read it.)"
            )

        messages = [
            SystemMessage(content=SYSTEM_PROMPT),
            HumanMessage(content=prompt),
        ]

        # Try the primary graph, then retry on transient/tool-validation errors.
        # tool-call validation failures consume an attempt (max 3); per-minute
        # rate limits just wait and retry without consuming one (max waits).
        result = None
        last_exc = None
        attempt = 0
        rate_waits = 0
        while attempt < 3:
            graph = self.graph if attempt == 0 else self._retry()
            try:
                result = graph.invoke(
                    {"messages": messages},
                    config={"recursion_limit": 25},
                )
                break
            except Exception as exc:  # noqa: BLE001 - keep the run alive
                last_exc = exc
                msg = str(exc).lower()
                # Per-minute rate limit: wait for the window to reset, then
                # retry the same graph without consuming a tool-fail attempt.
                if any(m in msg for m in _RATE_LIMIT_MARKERS) and rate_waits < 4:
                    rate_waits += 1
                    time.sleep(_RATE_LIMIT_WAIT)
                    continue
                # Retry tool-call validation failures at higher temperature;
                # bail on anything else (e.g. 402 quota) where retry won't help.
                if not any(m in msg for m in _TOOL_FAIL_MARKERS):
                    return f"AGENT ERROR: {exc}"
                attempt += 1
        if result is None:
            return f"AGENT ERROR: {last_exc}"

        convo = result["messages"]
        answer = _clean_answer(_message_text(convo[-1]))

        # Guard against degenerate finals: empty content, or the model echoing a
        # tool name instead of an answer (seen with smaller models). In that
        # case, force a final synthesis turn with no tools available.
        if not answer or answer.lower() in _TOOL_NAMES:
            try:
                synth = self.llm.invoke(
                    convo
                    + [
                        HumanMessage(
                            content=(
                                "Based on everything above, reply with ONLY the "
                                "final answer in the required format. No "
                                "explanation, no tool names."
                            )
                        )
                    ]
                )
                answer = _clean_answer(_message_text(synth))
            except Exception:  # noqa: BLE001 - keep whatever we had
                pass

        return answer


def _message_text(msg) -> str:
    """Extract plain text from a message whose content may be blocks."""
    content = getattr(msg, "content", msg)
    if isinstance(content, list):  # some providers return content blocks
        return " ".join(
            block.get("text", "") for block in content if isinstance(block, dict)
        )
    return content or ""