"""GAIA Level-1 agent for the HF Agents Course final assignment.

Built on smolagents CodeAgent. Works with ANY of these free LLM backends —
set whichever API key you can get and the agent auto-detects it:

    GROQ_API_KEY        console.groq.com/keys      (free, no card, fast)
    CEREBRAS_API_KEY    cloud.cerebras.ai          (free tier)
    OPENROUTER_API_KEY  openrouter.ai/keys         (has free models)
    MISTRAL_API_KEY     console.mistral.ai         (free tier)
    GOOGLE_API_KEY      aistudio.google.com/apikey (free, best multimodal)
    HF_TOKEN            huggingface.co/settings/tokens (needs credits)

Generic escape hatch for any other OpenAI-compatible endpoint:
    OPENAI_API_KEY + OPENAI_BASE_URL + AGENT_MODEL

Override the model with AGENT_MODEL if a default model id has been retired.
"""

import base64
import mimetypes
import os
import re
import tempfile
import time

import requests
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    tool,
)

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
WIKI_UA = "HF-Agents-Course-GAIA-Agent/1.0 (educational use)"

# Provider registry: env var -> (base_url, default model, label)
# Model ids are defaults only; override with AGENT_MODEL if one is retired.
PROVIDERS = [
    (
        "GROQ_API_KEY",
        "https://api.groq.com/openai/v1",
        "llama-3.3-70b-versatile",
        "Groq",
    ),
    (
        "CEREBRAS_API_KEY",
        "https://api.cerebras.ai/v1",
        "llama-3.3-70b",
        "Cerebras",
    ),
    (
        "OPENROUTER_API_KEY",
        "https://openrouter.ai/api/v1",
        "meta-llama/llama-3.3-70b-instruct:free",
        "OpenRouter",
    ),
    (
        "MISTRAL_API_KEY",
        "https://api.mistral.ai/v1",
        "mistral-large-latest",
        "Mistral",
    ),
    (
        "GOOGLE_API_KEY",
        "https://generativelanguage.googleapis.com/v1beta/openai/",
        "gemini-2.5-flash",
        "Gemini",
    ),
    (
        "OPENAI_API_KEY",
        os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
        "gpt-4o-mini",
        "OpenAI-compatible",
    ),
]


def _active_provider():
    """Return (api_key, base_url, model_id, label) for the first key found."""
    for env_var, base_url, default_model, label in PROVIDERS:
        key = os.getenv(env_var)
        if key:
            return key, base_url, os.getenv("AGENT_MODEL", default_model), label
    return None, None, None, None


API_KEY, BASE_URL, MODEL_ID, PROVIDER = _active_provider()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")


# --------------------------------------------------------------------------
# Gemini-only helper: native audio / video understanding
# --------------------------------------------------------------------------
def _gemini_generate(parts: list, retries: int = 3) -> str:
    model = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
    url = (
        "https://generativelanguage.googleapis.com/v1beta/models/"
        f"{model}:generateContent?key={GOOGLE_API_KEY}"
    )
    for attempt in range(retries):
        resp = requests.post(url, json={"contents": [{"parts": parts}]}, timeout=180)
        if resp.status_code == 429 and attempt < retries - 1:
            time.sleep(20 * (attempt + 1))
            continue
        resp.raise_for_status()
        return resp.json()["candidates"][0]["content"]["parts"][0]["text"]
    return "ERROR: Gemini rate limited."


def _inline_part(file_path: str) -> dict:
    mime = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
    with open(file_path, "rb") as f:
        return {
            "inline_data": {
                "mime_type": mime,
                "data": base64.b64encode(f.read()).decode(),
            }
        }


# --------------------------------------------------------------------------
# Tools
# --------------------------------------------------------------------------
@tool
def wikipedia_page(title: str) -> str:
    """Fetch the full plain text of an English Wikipedia article. Use this
    instead of visit_webpage for Wikipedia — it never gets blocked.

    Args:
        title: Article title, e.g. "Mercedes Sosa" or "1928 Summer Olympics".
    """
    try:
        resp = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={
                "action": "query",
                "prop": "extracts",
                "explaintext": 1,
                "redirects": 1,
                "format": "json",
                "titles": title,
            },
            headers={"User-Agent": WIKI_UA},
            timeout=45,
        )
        resp.raise_for_status()
        pages = resp.json()["query"]["pages"]
        page = list(pages.values())[0]
        if "extract" not in page:
            return f"No Wikipedia article found for '{title}'. Try wikipedia_search."
        return page["extract"][:60000]
    except Exception as e:
        return f"ERROR fetching Wikipedia page: {e}"


@tool
def wikipedia_search(query: str) -> str:
    """Search English Wikipedia and return matching article titles with snippets.
    Use this to find the right title, then call wikipedia_page.

    Args:
        query: Search terms.
    """
    try:
        resp = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={
                "action": "query",
                "list": "search",
                "srsearch": query,
                "srlimit": 10,
                "format": "json",
            },
            headers={"User-Agent": WIKI_UA},
            timeout=45,
        )
        resp.raise_for_status()
        hits = resp.json()["query"]["search"]
        return "\n".join(
            f"- {h['title']}: {re.sub('<[^<]+?>', '', h['snippet'])}" for h in hits
        ) or "No results."
    except Exception as e:
        return f"ERROR searching Wikipedia: {e}"


@tool
def fetch_url(url: str) -> str:
    """Fetch a web page as text with a browser-like user agent. Use when
    visit_webpage fails with a 403 Forbidden error.

    Args:
        url: The full URL to fetch.
    """
    try:
        resp = requests.get(
            url,
            headers={
                "User-Agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/122.0 Safari/537.36"
                )
            },
            timeout=60,
        )
        resp.raise_for_status()
        try:
            from markdownify import markdownify

            text = markdownify(resp.text)
        except Exception:
            text = resp.text
        text = re.sub(r"\n{3,}", "\n\n", text)
        return text[:50000]
    except Exception as e:
        return f"ERROR fetching url: {e}"


@tool
def transcribe_audio(file_path: str) -> str:
    """Transcribe a local audio file (mp3/wav/m4a) to English text.

    Args:
        file_path: Absolute path to the local audio file to transcribe.
    """
    # Gemini: native audio understanding
    if GOOGLE_API_KEY:
        try:
            return _gemini_generate(
                [{"text": "Transcribe this audio verbatim."}, _inline_part(file_path)]
            )
        except Exception as e:
            return f"ERROR transcribing audio: {e}"
    # Groq hosts Whisper on an OpenAI-compatible endpoint
    if os.getenv("GROQ_API_KEY"):
        try:
            with open(file_path, "rb") as f:
                resp = requests.post(
                    "https://api.groq.com/openai/v1/audio/transcriptions",
                    headers={"Authorization": f"Bearer {os.getenv('GROQ_API_KEY')}"},
                    files={"file": (os.path.basename(file_path), f)},
                    data={"model": os.getenv("ASR_MODEL", "whisper-large-v3")},
                    timeout=180,
                )
            resp.raise_for_status()
            return resp.json()["text"]
        except Exception as e:
            return f"ERROR transcribing audio: {e}"
    # HF Inference fallback
    try:
        from huggingface_hub import InferenceClient

        result = InferenceClient(
            token=os.getenv("HF_TOKEN")
        ).automatic_speech_recognition(
            file_path, model=os.getenv("ASR_MODEL", "openai/whisper-large-v3")
        )
        return result.text if hasattr(result, "text") else str(result)
    except Exception as e:
        return f"ERROR transcribing audio (no ASR backend available): {e}"


@tool
def analyze_image(file_path: str, question: str) -> str:
    """Answer a question about a local image file using a vision model.

    Args:
        file_path: Absolute path to the local image file (png/jpg).
        question: The question to answer about the image. Be specific; for
            chess positions, ask for a full square-by-square board reading
            AND the winning move, verified carefully.
    """
    if GOOGLE_API_KEY:
        try:
            return _gemini_generate([{"text": question}, _inline_part(file_path)])
        except Exception as e:
            return f"ERROR analyzing image: {e}"
    if not API_KEY:
        return "ERROR: no vision backend configured."
    try:
        mime = mimetypes.guess_type(file_path)[0] or "image/png"
        with open(file_path, "rb") as f:
            b64 = base64.b64encode(f.read()).decode()
        vision_model = os.getenv("VISION_MODEL", MODEL_ID)
        resp = requests.post(
            BASE_URL.rstrip("/") + "/chat/completions",
            headers={"Authorization": f"Bearer {API_KEY}"},
            json={
                "model": vision_model,
                "max_tokens": 1500,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": question},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:{mime};base64,{b64}"},
                            },
                        ],
                    }
                ],
            },
            timeout=180,
        )
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"]
    except Exception as e:
        return (
            f"ERROR analyzing image: {e}. The model may not support images; "
            "set VISION_MODEL to a vision-capable model id."
        )


@tool
def analyze_youtube_video(video_url: str, question: str) -> str:
    """Watch a YouTube video and answer a question about its visual and audio
    content (counting things on screen, quotes, scenes). Requires a Gemini key.

    Args:
        video_url: Full YouTube URL, e.g. https://www.youtube.com/watch?v=XXXX
        question: The question to answer about the video.
    """
    if not GOOGLE_API_KEY:
        return (
            "ERROR: video analysis needs GOOGLE_API_KEY. Use "
            "get_youtube_transcript or web_search for descriptions instead."
        )
    try:
        return _gemini_generate(
            [{"text": question}, {"file_data": {"file_uri": video_url}}]
        )
    except Exception as e:
        return f"ERROR analyzing video: {e}"


@tool
def get_youtube_transcript(video_url: str) -> str:
    """Fetch the transcript/captions of a YouTube video as plain text.

    Args:
        video_url: Full YouTube URL, e.g. https://www.youtube.com/watch?v=XXXX
    """
    try:
        from youtube_transcript_api import YouTubeTranscriptApi

        m = re.search(r"(?:v=|youtu\.be/)([\w-]{11})", video_url)
        if not m:
            return "ERROR: could not parse video id from URL."
        vid = m.group(1)
        try:
            entries = [s.text for s in YouTubeTranscriptApi().fetch(vid)]
        except AttributeError:
            entries = [s["text"] for s in YouTubeTranscriptApi.get_transcript(vid)]
        return " ".join(entries)[:20000]
    except Exception as e:
        return f"ERROR fetching transcript: {e}. Try web_search instead."


@tool
def read_file_as_text(file_path: str) -> str:
    """Read a local text-like file (py, txt, csv, json, md) and return its content.

    Args:
        file_path: Absolute path to the local file.
    """
    try:
        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            return f.read()[:30000]
    except Exception as e:
        return f"ERROR reading file: {e}"


# --------------------------------------------------------------------------
# Answer-format guidance (GAIA is scored by EXACT MATCH)
# --------------------------------------------------------------------------
GAIA_INSTRUCTIONS = """You are a general AI assistant answering a benchmark
question scored by EXACT string match. Work step by step with your tools,
then call final_answer() with ONLY the answer itself.

Formatting rules for the final answer (critical):
- Do NOT write "FINAL ANSWER" or any prefix/suffix, explanation, or period
  at the end. Output the bare answer only.
- Numbers: plain digits, no thousands separators, no units ($, %, kg) unless
  the question explicitly asks for them, no trailing ".0".
- Strings: no articles ("the", "a"), no abbreviations unless asked.
- Comma-separated lists: apply the rules above to each element, use ", "
  (comma + space) between elements, and respect any ordering the question
  asks for (e.g. alphabetical, ascending).
- If asked for a first name / last name / city / country code only, return
  exactly that and nothing more.

Tool strategy:
- Wikipedia questions: use wikipedia_search then wikipedia_page. Do NOT use
  visit_webpage on wikipedia.org — it returns 403. The article text often
  contains a discography or results table; read it carefully and count.
- If visit_webpage returns 403 Forbidden, retry that URL with fetch_url.
- Attached files: a local path is given; use read_file_as_text,
  transcribe_audio, analyze_image, or pandas (pd.read_excel) for .xlsx.
- For .xlsx, inspect the columns first, then compute. Format money like
  89706.00 only when the question asks for two decimal places.
- Python-code questions: read the code and reason through it carefully.
- YouTube: try analyze_youtube_video, then get_youtube_transcript, then
  web_search for third-party descriptions of the video.
- Some questions are pure reasoning (reversed text, a group-theory table).
  Solve those directly in python without searching.

Reliability rules:
- Never give up and guess a number you did not verify. If one source is
  blocked, try another tool or another source.
- Re-read the question's exact wording before answering (e.g. "included",
  "as of July 2023", "IOC country code", "without abbreviations",
  "first name only").
"""


class GAIAAgent:
    """Wraps a smolagents CodeAgent with GAIA-specific tooling and prompting."""

    def __init__(self):
        if not API_KEY:
            raise RuntimeError(
                "No LLM API key found. Set one of: GROQ_API_KEY, "
                "CEREBRAS_API_KEY, OPENROUTER_API_KEY, MISTRAL_API_KEY, "
                "GOOGLE_API_KEY, or OPENAI_API_KEY (+OPENAI_BASE_URL)."
            )
        from smolagents import OpenAIServerModel

        model = OpenAIServerModel(
            model_id=MODEL_ID, api_base=BASE_URL, api_key=API_KEY
        )
        self.agent = CodeAgent(
            model=model,
            tools=[
                DuckDuckGoSearchTool(),
                VisitWebpageTool(),
                fetch_url,
                wikipedia_search,
                wikipedia_page,
                transcribe_audio,
                analyze_image,
                analyze_youtube_video,
                get_youtube_transcript,
                read_file_as_text,
            ],
            additional_authorized_imports=[
                "pandas",
                "numpy",
                "openpyxl",
                "json",
                "csv",
                "re",
                "math",
                "statistics",
                "itertools",
                "collections",
                "datetime",
            ],
            max_steps=15,
        )
        print(f"GAIAAgent initialized (backend={PROVIDER}, model={MODEL_ID}).")

    @staticmethod
    def download_task_file(task_id: str, file_name: str) -> str | None:
        """Download the file attached to a task; returns a local path or None."""
        if not file_name:
            return None
        url = f"{DEFAULT_API_URL}/files/{task_id}"
        for attempt in range(4):
            try:
                resp = requests.get(url, timeout=30)
                resp.raise_for_status()
                fd, path = tempfile.mkstemp(suffix=os.path.splitext(file_name)[1] or "")
                with os.fdopen(fd, "wb") as f:
                    f.write(resp.content)
                return path
            except Exception as e:
                print(f"File download attempt {attempt + 1} failed for {task_id}: {e}")
                time.sleep(3 * (attempt + 1))
        return None

    @staticmethod
    def _clean(answer: str) -> str:
        """Strip wrappers the model sometimes adds despite instructions."""
        a = str(answer).strip()
        a = re.sub(r"^(final answer\s*:?\s*)", "", a, flags=re.IGNORECASE)
        a = a.strip().strip('"').strip("'").strip()
        if a.endswith("."):
            a = a[:-1]
        return a

    def __call__(self, question: str, task_id: str = "", file_name: str = "") -> str:
        prompt = GAIA_INSTRUCTIONS + "\n\nQuestion: " + question
        file_path = self.download_task_file(task_id, file_name)
        if file_path:
            prompt += (
                f"\n\nAn attached file for this question was downloaded to the "
                f"local path: {file_path} (original name: {file_name})."
            )
        elif file_name:
            prompt += (
                f"\n\nNOTE: this question references an attached file "
                f"({file_name}) that could not be downloaded. Answer from "
                f"other sources if possible."
            )
        try:
            return self._clean(self.agent.run(prompt))
        except Exception as e:
            print(f"Agent error on task {task_id}: {e}")
            return f"AGENT ERROR: {e}"