from __future__ import annotations

import os
import re

from smolagents import (
    ToolCallingAgent,
    OpenAIModel,
    PythonInterpreterTool,
    FinalAnswerTool,
)

from gaia_tools import build_tools
import gemini_limits as limits
import providers
import vertex_auth


AGENT_INSTRUCTIONS = r"""
You are a high-reliability solver for the Hugging Face Agents Course Unit 4
GAIA benchmark (Level 1 questions). Answers are scored by EXACT MATCH, so the
final answer must be short, precise and in the format the question asks for.

GENERAL RULES
1. Read the entire question first. Note exactly what is asked (a number, a
   name, a list, a single word, a move), any range (inclusive?), order, units
   and formatting requirements.
2. Pick the right tool for the job (playbook below). Never repeat a call that
   already failed or returned nothing: change the query, the tool or the source.
3. Do not stop at the first plausible answer: sanity-check it against the
   question (right entity, right year, right range, right unit).
4. Do arithmetic, counting, sorting and filtering in python_interpreter, not in
   your head. List the items explicitly before counting them.
5. You have a limited number of steps. Be efficient: one good tool call beats
   three vague ones.

PLAYBOOK BY QUESTION TYPE
- English Wikipedia facts, lists, discographies, tables, statistics: call
  wikipedia_page(title, year?, section?). Use year=2022 when the question says
  "2022 version"/"as of 2022"; use section="Discography" etc. to jump to a table.
  Wikipedia project pages also work (e.g. "Wikipedia:Featured articles promoted
  in 2016", "Wikipedia:Featured article candidates/<Title>/archive1"). Use
  wikipedia_search to find the exact title.
- Hard web lookups (papers, awards, grant numbers, rosters, obscure facts), or
  when web_search returns irrelevant pages: use google_search with a specific
  question, then open a cited source with visit_webpage to confirm.
  web_search: plain keywords (no site: operators); visit_webpage: any URL.
- YouTube: youtube_transcript for spoken words; analyze_youtube_video for
  anything visual (counting things on screen) or when no transcript exists, or
  when you must hear the audio. Never web-search a video ID.
- Attachments (the task lists a file): ALWAYS call inspect_attachment first.
  * image: read the extracted text. CHESS: take the FEN it prints, call
    chess_best_move, and answer with the move in algebraic notation for the
    correct side to move (double-check the FEN piece count and orientation).
  * audio (.mp3): inspect_attachment returns a transcript; extract exactly what
    is asked (e.g. only ingredients, sorted, without quantities).
  * .py: call run_attached_python to get the real output; if it is random or
    fails, reason from the source.
  * .xlsx/.csv: the dump shows the data and its local path; compute with
    python_interpreter (pandas is available: pd.read_excel(path)); respect
    exclusions in the question (e.g. "not including drinks").
- Text that looks reversed, scrambled or encoded: decode it with Python first
  (e.g. text[::-1]), then answer the decoded question literally (if it asks for
  the opposite of a word, answer just that word).
- Tables or rules given in the question (e.g. an operation table): copy them
  into Python and check every case by brute force.
- Classification questions (e.g. botany): apply the strict definition the
  question implies (botanical fruit vs. culinary vegetable), go item by item,
  then format as asked.
- python_interpreter has NO internet; use web tools for lookups.

ANSWER FORMAT (exact match!)
- Call final_answer with ONLY the answer: no "Final Answer:", no explanation,
  no citations, no Markdown, no quotation marks, no trailing period.
- A number: digits only (no thousands separators, no units) unless the question
  asks for units or a specific format.
- A list: comma-separated, in the order/sorting the question asks (alphabetical
  if asked), no brackets, no "and".
- Names: give exactly the part asked (e.g. only the last name, or the full
  name). Use the spelling used by the source. Words in lowercase unless they
  are proper names.
- If after real effort you cannot verify the answer, still give your single
  best answer (never say "unknown" or explain).
"""


def _clean_answer(value: object) -> str:
    if value is None:
        return ""

    text = str(value).strip()

    # Remove accidental Markdown fences.
    if text.startswith("```") and text.endswith("```"):
        lines = text.splitlines()
        if len(lines) >= 2:
            text = "\n".join(lines[1:-1]).strip()

    # Remove common wrapper labels.
    text = re.sub(
        r"^\s*(FINAL ANSWER|Final Answer|Answer|Final)\s*:\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )

    text = text.strip().strip('"').strip("'")

    # Step-limit / tool-call leftovers are not answers.
    if text.lower().startswith(("calling tools:", "error:", "agent_error:")):
        return ""

    # A single trailing period on a short answer ("Right.") breaks exact match.
    if text.endswith(".") and text.count(".") == 1 and len(text.split()) <= 8:
        text = text[:-1].rstrip()
    return text


class RouterModel(OpenAIModel):
    """One model object that fails over across providers AND models.

    Before every request it picks the first endpoint (provider + model) with
    spare per-minute budget. On 429 / daily-quota / rejected key / provider
    quirks the endpoint is parked or dropped and the SAME conversation simply
    continues on the next one, instead of restarting the question.
    """

    def __init__(self, endpoints, **kwargs):
        self.endpoints = {e.key: e for e in endpoints}
        self._clients = {}
        self._last_key = None
        first = endpoints[0]
        # retry=False: smolagents' built-in retry just sleeps and hammers the
        # same rate-limited model; failover is handled here.
        super().__init__(
            model_id=first.model, api_base=first.base_url,
            api_key=first.api_key, retry=False, **kwargs,
        )

    def _client_for(self, endpoint):
        if endpoint.provider == "gemini":
            # Vertex tokens expire; obtain a fresh ADC token for every request.
            return vertex_auth.openai_client()

        import openai

        if endpoint.provider not in self._clients:
            self._clients[endpoint.provider] = openai.OpenAI(
                api_key=endpoint.api_key, base_url=endpoint.base_url,
                max_retries=0,  # the SDK's hidden sleep-and-retry is unwanted
                timeout=120,
            )
        return self._clients[endpoint.provider]

    def generate(self, messages, *args, **kwargs):
        # Providers with tiny per-minute token budgets (e.g. Groq free) are
        # skipped for big prompts instead of failing with 413/429.
        est_tokens = len(str(messages)) / 3.5
        keys = [
            k for k, e in self.endpoints.items()
            if not (e.max_prompt_tokens and est_tokens > e.max_prompt_tokens)
        ] or list(self.endpoints)

        last_exc = None
        for _ in range(len(keys) * 6):
            key = limits.acquire(keys)  # may raise limits.QuotaExhausted
            ep = self.endpoints[key]
            self.client = self._client_for(ep)
            self.model_id = ep.model
            # Non-Gemini providers are happiest with plain-string content.
            self.flatten_messages_as_text = ep.provider != "gemini"
            if key != self._last_key:
                print(f"[LLM] using {key}")
                self._last_key = key
            try:
                return super().generate(messages, *args, **kwargs)
            except Exception as exc:
                last_exc = exc
                if not limits.report_error(key, exc):
                    raise
        raise last_exc


def _log_step(step, agent=None):
    """Print what each step did so failures are visible (verbosity is off)."""
    try:
        calls = ", ".join(
            f"{tc.name}({str(tc.arguments)[:120]})"
            for tc in (getattr(step, "tool_calls", None) or [])
        )
        obs = str(getattr(step, "observations", "") or "").replace("\n", " ")
        err = getattr(step, "error", None)
        print(f"[STEP {step.step_number}] {calls or '(no tool call)'}")
        if err:
            print(f"    ERROR: {str(err)[:200]}")
        elif obs:
            print(f"    -> {obs[:200]}")
    except Exception:
        pass


def _find_quota_error(exc):
    seen = set()
    while exc is not None and id(exc) not in seen:
        if isinstance(exc, limits.QuotaExhausted):
            return exc
        seen.add(id(exc))
        exc = exc.__cause__ or exc.__context__
    return None


class GaiaAgent:
    def __init__(self, task_map):
        self.task_map = task_map
        # Raises a clear error if no provider key is configured.
        self.endpoints = providers.build_endpoints()
        names = ", ".join(dict.fromkeys(e.provider for e in self.endpoints))
        print(f"[LLM] providers in order: {names} "
              f"({len(self.endpoints)} model endpoints)")

    def _build_agent(self):
        model = RouterModel(
            self.endpoints,
            temperature=0.0,
            max_tokens=2500,
        )

        tools = build_tools(self.task_map)
        tools.append(
            PythonInterpreterTool(
                authorized_imports=[
                    "math", "statistics", "datetime", "json", "re", "csv",
                    "collections", "itertools", "fractions", "decimal",
                    "string", "unicodedata", "random", "time",
                    "pandas", "pandas.*", "numpy", "numpy.*", "openpyxl",
                    "openpyxl.*",
                ]
            )
        )
        tools.append(FinalAnswerTool())

        return ToolCallingAgent(
            tools=tools,
            model=model,
            instructions=AGENT_INSTRUCTIONS,
            max_steps=int(os.getenv("MAX_STEPS", "12")),
            planning_interval=None,
            verbosity_level=0,
            step_callbacks=[_log_step],
        )

    def solve(self, question, task_id, file_name=""):
        prompt = f"""
GAIA TASK ID:
{task_id}

ATTACHMENT:
{file_name or "None"}

USER QUESTION:
{question}

Solve this task accurately.

Use tools when needed.
Verify important facts.
Follow the requested output format exactly.

When finished, return ONLY the answer.
Do not include explanations, citations, Markdown, or "Final Answer:".
"""
        agent = self._build_agent()
        try:
            result = agent.run(prompt)
        except Exception as exc:
            # smolagents wraps model errors (AgentGenerationError), so look
            # through the exception chain for our own quota error.
            quota = _find_quota_error(exc)
            if quota is not None:
                print(f"[AGENT] out of quota: {quota}")
                return f"AGENT_ERROR: QuotaExhausted: {quota}"
            print(f"[AGENT] failed: {type(exc).__name__}: {exc}")
            return f"AGENT_ERROR: {type(exc).__name__}: {exc}"

        answer = _clean_answer(result)
        if answer:
            return answer

        # Ran out of steps: ask once, without tools, for a best answer based
        # on everything gathered so far.
        try:
            messages = agent.write_memory_to_messages()
            messages.append({
                "role": "user",
                "content": [{"type": "text", "text": (
                    "You are out of steps. Using only what you found so far "
                    "(and your own knowledge if needed), reply with your best "
                    "final answer ONLY: no explanation, no tool calls."
                )}],
            })
            reply = agent.model.generate(messages)
            answer = _clean_answer(reply.content)
            if answer:
                print("[AGENT] step limit reached; used best-effort answer.")
                return answer
        except Exception as exc:
            print(f"[AGENT] rescue failed: {type(exc).__name__}: {exc}")

        return "AGENT_ERROR: no answer produced (step limit reached or tool failures)"