"""
GAIA agent for the Hugging Face Agents Course (Unit 4 final assignment).

Default: local Ollama (same as your original setup). Configure with env vars:

  LLM_BACKEND   "ollama" (default) | "hf" (HF Inference, needs HF_TOKEN) | "litellm"
  LLM_MODEL_ID  default for ollama: ollama_chat/qwen2.5:7b
  LLM_API_BASE  default for ollama: http://localhost:11434
  LLM_API_KEY   only for "litellm"
  NUM_CTX       Ollama context window (default 8192)
  MAX_PAGE_CHARS  max characters a web tool returns (default 5000, keeps the 8k context safe)
"""

import os
import re
import tempfile
from pathlib import Path
from urllib.parse import urlparse, parse_qs

import pandas as pd
import requests
from bs4 import BeautifulSoup
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    InferenceClientModel,
    LiteLLMModel,
    tool,
)

MAX_PAGE_CHARS = int(os.getenv("MAX_PAGE_CHARS", "5000"))
HEADERS = {"User-Agent": "Mozilla/5.0 (GAIA course agent; educational use)"}


# --------------------------------------------------------------------------- #
# Model
# --------------------------------------------------------------------------- #
def build_model():
    backend = os.getenv("LLM_BACKEND", "ollama").lower()

    if backend == "ollama":
        return LiteLLMModel(
            model_id=os.getenv("LLM_MODEL_ID", "ollama_chat/qwen2.5:7b"),
            api_base=os.getenv("LLM_API_BASE", "http://localhost:11434"),
            api_key="ollama",
            num_ctx=int(os.getenv("NUM_CTX", "8192")),
        )

    if backend == "litellm":
        kwargs = {"model_id": os.environ["LLM_MODEL_ID"]}
        if os.getenv("LLM_API_KEY"):
            kwargs["api_key"] = os.getenv("LLM_API_KEY")
        if os.getenv("LLM_API_BASE"):
            kwargs["api_base"] = os.getenv("LLM_API_BASE")
        return LiteLLMModel(**kwargs)

    token = os.getenv("HF_TOKEN")
    if not token:
        raise RuntimeError("HF_TOKEN is not set (needed for LLM_BACKEND=hf).")
    return InferenceClientModel(
        model_id=os.getenv("LLM_MODEL_ID", "Qwen/Qwen2.5-72B-Instruct"), token=token
    )


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _page_text(html: str) -> str:
    soup = BeautifulSoup(html, "html.parser")
    for tag in soup(["script", "style", "nav", "footer", "header", "aside"]):
        tag.decompose()
    return soup.get_text(separator="\n", strip=True)


def download_task_file(api_url: str, task_id: str, file_name: str | None = None) -> str | None:
    """Download the file attached to a task (if any). Returns a local path or None."""
    try:
        r = requests.get(f"{api_url}/files/{task_id}", timeout=30)
        if r.status_code == 404:
            return None
        r.raise_for_status()
    except requests.RequestException as e:
        print(f"Could not download file for task {task_id}: {e}")
        return None

    if not file_name:
        cd = r.headers.get("content-disposition", "")
        m = re.search(r'filename="?([^";]+)"?', cd)
        file_name = m.group(1) if m else task_id

    out_dir = Path(tempfile.gettempdir()) / "gaia_files" / task_id
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / Path(file_name).name
    path.write_bytes(r.content)
    return str(path)


def clean_answer(text: str) -> str:
    """The scorer is EXACT MATCH, so strip anything that is not the answer."""
    t = str(text).strip()
    # Never submit the literal "FINAL ANSWER" marker (course requirement).
    t = re.sub(r"(?is)^.*?final answer\s*[:\-]\s*", "", t).strip()
    # Remove trailing sentence periods, then wrapping quotes.
    t = t.rstrip(".").strip()
    if len(t) >= 2 and t[0] == t[-1] and t[0] in "\"'":
        t = t[1:-1].strip()
    return t


# --------------------------------------------------------------------------- #
# Tools
# --------------------------------------------------------------------------- #
@tool
def visit_webpage(url: str) -> str:
    """
    Read the textual content of a public webpage.

    Args:
        url: The URL of the webpage to retrieve.
    """
    try:
        r = requests.get(url, timeout=20, headers=HEADERS)
        r.raise_for_status()
    except requests.RequestException as e:
        return f"Error fetching {url}: {e}"
    return _page_text(r.text)[:MAX_PAGE_CHARS]


@tool
def extract_webpage_section(url: str, keyword: str) -> str:
    """
    Retrieve a webpage and return the text around a keyword (useful for long pages).

    Args:
        url: URL of the webpage.
        keyword: Section title or phrase to search for.
    """
    try:
        r = requests.get(url, timeout=20, headers=HEADERS)
        r.raise_for_status()
    except requests.RequestException as e:
        return f"Error fetching {url}: {e}"

    text = _page_text(r.text)
    pos = text.lower().find(keyword.lower())
    if pos == -1:
        return "Keyword not found."
    start = max(0, pos - 300)
    return text[start : start + MAX_PAGE_CHARS]


@tool
def wikipedia_search(query: str) -> str:
    """
    Search English Wikipedia and return the titles and snippets of the best matches.

    Args:
        query: What to search for.
    """
    try:
        r = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={
                "action": "query",
                "list": "search",
                "srsearch": query,
                "srlimit": 5,
                "format": "json",
            },
            headers=HEADERS,
            timeout=20,
        )
        r.raise_for_status()
        hits = r.json()["query"]["search"]
    except Exception as e:
        return f"Wikipedia search failed: {e}"
    if not hits:
        return "No results."
    return "\n".join(
        f"- {h['title']}: {BeautifulSoup(h['snippet'], 'html.parser').get_text()}"
        for h in hits
    )


@tool
def wikipedia_page(title: str, section_keyword: str = "") -> str:
    """
    Get the plain text of an English Wikipedia article (tables included as text).

    Args:
        title: Exact article title, e.g. "Mercedes Sosa".
        section_keyword: Optional word/phrase (e.g. "Discography") - if given, return the text around it.
    """
    try:
        r = requests.get(
            f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}",
            headers=HEADERS,
            timeout=20,
        )
        r.raise_for_status()
    except requests.RequestException as e:
        return f"Error fetching article '{title}': {e}"

    soup = BeautifulSoup(r.text, "html.parser")
    content = soup.find(id="mw-content-text") or soup
    for tag in content(["script", "style", "sup"]) + content.select(".navbox, .reflist"):
        tag.decompose()
    text = content.get_text(separator="\n", strip=True)

    if section_keyword:
        pos = text.lower().find(section_keyword.lower())
        if pos != -1:
            return text[max(0, pos - 200) : pos + MAX_PAGE_CHARS]
    return text[:MAX_PAGE_CHARS]


@tool
def read_file(path: str) -> str:
    """
    Read a file attached to the task. Supports text/code/json/csv/xlsx files.
    For audio use transcribe_audio, for images use describe_image.

    Args:
        path: Local path of the file.
    """
    p = Path(path)
    if not p.exists():
        return f"File not found: {path}"
    ext = p.suffix.lower()
    try:
        if ext in (".xlsx", ".xls"):
            sheets = pd.read_excel(p, sheet_name=None)
            return "\n\n".join(
                f"## Sheet: {name}\n{df.to_string()}" for name, df in sheets.items()
            )[:20000]
        if ext == ".csv":
            return pd.read_csv(p).to_string()[:20000]
        return p.read_text(errors="replace")[:20000]
    except Exception as e:
        return f"Could not read {path}: {e}"


@tool
def transcribe_audio(path: str) -> str:
    """
    Transcribe a local audio file (mp3, wav, ...) to text.

    Args:
        path: Local path of the audio file.
    """
    try:
        from huggingface_hub import InferenceClient

        client = InferenceClient(token=os.getenv("HF_TOKEN"))
        out = client.automatic_speech_recognition(
            path, model=os.getenv("ASR_MODEL_ID", "openai/whisper-large-v3")
        )
        return out.text
    except Exception as e:
        return f"Transcription failed: {e}"


@tool
def describe_image(path: str, question: str) -> str:
    """
    Ask a vision model a question about a local image file.

    Args:
        path: Local path of the image.
        question: What you want to know about the image.
    """
    try:
        import base64
        import mimetypes
        from huggingface_hub import InferenceClient

        mime = mimetypes.guess_type(path)[0] or "image/png"
        b64 = base64.b64encode(Path(path).read_bytes()).decode()
        client = InferenceClient(token=os.getenv("HF_TOKEN"))
        resp = client.chat_completion(
            model=os.getenv("VISION_MODEL_ID", "Qwen/Qwen2.5-VL-72B-Instruct"),
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}},
                        {"type": "text", "text": question},
                    ],
                }
            ],
            max_tokens=500,
        )
        return resp.choices[0].message.content
    except Exception as e:
        return f"Image analysis failed: {e}"


@tool
def youtube_transcript(url: str) -> str:
    """
    Get the transcript (captions) of a YouTube video.

    Args:
        url: The YouTube video URL.
    """
    try:
        from youtube_transcript_api import YouTubeTranscriptApi

        parsed = urlparse(url)
        video_id = parse_qs(parsed.query).get("v", [None])[0] or parsed.path.strip("/").split("/")[-1]
        segments = YouTubeTranscriptApi().fetch(video_id)
        return " ".join(s.text for s in segments)[:MAX_PAGE_CHARS * 2]
    except Exception as e:
        return f"Could not get transcript: {e}"


# --------------------------------------------------------------------------- #
# Agent
# --------------------------------------------------------------------------- #
INSTRUCTIONS = """
You are solving questions from the GAIA benchmark. Answers are graded by EXACT MATCH.

Method:
- Work out exactly what is asked and which source/date/category it specifies; use that source.
- Use the tools: web search, Wikipedia tools, visit_webpage, and for attached files
  read_file / transcribe_audio / describe_image. Use python for any counting or arithmetic
  instead of doing it in your head.
- For counting or list questions, enumerate every candidate first, check each against ALL
  constraints (e.g. studio albums only, date ranges inclusive/exclusive), then count.
- Never invent facts. If one source fails, try another.

Final answer format (very important):
- Call final_answer(...) with ONLY the answer - no explanation, no sentence, no "FINAL ANSWER".
- Numbers: plain digits, no thousands separators, no units or % unless the question asks.
- Strings: no articles (a/an/the) and no abbreviations unless asked; digits written as digits.
- Lists: comma-separated, in the order requested, following the same rules per element.
""".strip()


class BasicAgent:
    def __init__(self):
        self.model = build_model()
        self.agent = CodeAgent(
            tools=[
                DuckDuckGoSearchTool(max_results=5),
                wikipedia_search,
                wikipedia_page,
                visit_webpage,
                extract_webpage_section,
                read_file,
                transcribe_audio,
                describe_image,
                youtube_transcript,
            ],
            model=self.model,
            instructions=INSTRUCTIONS,
            additional_authorized_imports=[
                "re", "math", "statistics", "itertools", "collections",
                "datetime", "json", "csv", "pandas", "numpy", "unicodedata",
            ],
            max_steps=10,
            verbosity_level=1,
        )

    def __call__(self, question: str, file_path: str | None = None) -> str:
        task = question
        if file_path:
            task += (
                f"\n\nAn attached file for this question is available locally at: {file_path}"
                "\nOpen it with the appropriate tool (read_file, transcribe_audio, describe_image)."
            )

        last_error = None
        for _ in range(2):  # one retry for transient API / search failures
            try:
                return clean_answer(self.agent.run(task))
            except Exception as e:  # noqa: BLE001
                last_error = e
                print(f"Agent attempt failed: {e}")
        raise last_error