"""GAIA agent for the HF Agents course, Unit 4. smolagents CodeAgent + free Gemini API."""
import mimetypes
import os
import re

from dotenv import load_dotenv
from google import genai
from google.genai import types
import requests
from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel, WikipediaSearchTool, tool

load_dotenv()
# Exact model names, tried in order: the free daily quota is per model, so when one runs out the next takes over.
GEMINI_MODELS = os.environ.get(
    "GEMINI_MODELS",
    "gemini-3-flash-preview,gemini-3.5-flash,gemini-3.6-flash,gemini-3.7-flash,gemini-3.8-flash,gemini-3.5-flash-lite,gemini-3.1-flash-lite",
).split(",")
USER_AGENT = "GAIA-course-agent/1.0 (Hugging Face Agents course; https://huggingface.co/brunopaschoal)"  # Wikipedia 403s anonymous bots
_gemini = genai.Client(api_key=os.environ["GEMINI_API_KEY"])


def _ask_gemini(parts: list) -> str:
    for model in GEMINI_MODELS[:-1]:
        try:
            return _gemini.models.generate_content(model=model, contents=parts).text or ""
        except Exception:
            continue  # quota or model error: try the next model
    return _gemini.models.generate_content(model=GEMINI_MODELS[-1], contents=parts).text or ""


@tool
def visit_webpage(url: str) -> str:
    """Open a web page (e.g. a Wikipedia article) and return its content as markdown text.
    Args:
        url: The full URL of the page.
    """
    from markdownify import markdownify

    try:
        r = requests.get(url, headers={"User-Agent": USER_AGENT}, timeout=30)
        r.raise_for_status()
    except Exception as e:
        return f"Error fetching {url}: {e}"
    text = re.sub(r"\n{3,}", "\n\n", markdownify(r.text))
    return text[:40000] + ("\n...[page truncated]" if len(text) > 40000 else "")  # keeps tokens/minute under the free limit


@tool
def youtube_transcript(url: str) -> str:
    """Get the spoken words (captions) of a YouTube video as text. Cheaper than ask_about_youtube; use it first for questions about what is said.
    Args:
        url: The full YouTube URL, e.g. 'https://www.youtube.com/watch?v=...'.
    """
    from youtube_transcript_api import YouTubeTranscriptApi

    video_id = re.search(r"(?:v=|youtu\.be/)([\w-]{11})", url).group(1)
    return " ".join(s.text for s in YouTubeTranscriptApi().fetch(video_id))


@tool
def ask_about_file(file_path: str, question: str) -> str:
    """Ask a question about an image or audio file (e.g. a chess board picture, or an mp3 recording to transcribe).
    Args:
        file_path: Local path of the image or audio file.
        question: What you want to know about the file. Be specific, e.g. 'Transcribe this audio word for word.'
    """
    mime = mimetypes.guess_type(file_path)[0] or "application/octet-stream"
    with open(file_path, "rb") as f:
        return _ask_gemini([types.Part.from_bytes(data=f.read(), mime_type=mime), question])


@tool
def ask_about_youtube(url: str, question: str) -> str:
    """Watch a YouTube video (images and sound) and answer a question about it.
    Args:
        url: The full YouTube URL, e.g. 'https://www.youtube.com/watch?v=...'.
        question: What you want to know about the video.
    """
    return _ask_gemini([types.Part(file_data=types.FileData(file_uri=url)), question])


INSTRUCTIONS = """You are answering GAIA benchmark questions, scored by EXACT MATCH against the true answer.
- Research carefully with the tools: search the web, read Wikipedia and web pages, and double-check facts.
- For Wikipedia facts, open the article directly with visit_webpage("https://en.wikipedia.org/wiki/<Title>").
  If web_search fails or is rate limited, use wikipedia_search or visit_webpage instead of retrying the same search.
- Use ask_about_file for images/audio. For YouTube: youtube_transcript for what is said, ask_about_youtube for what is shown.
- For an attached Excel file, load it with pandas.read_excel(path) and compute the answer in code.
- Write ONE short code block per step, then stop and wait for its output. Never invent tool outputs.
- Call final_answer with ONLY the answer: no explanation, no "FINAL ANSWER" text, no sentence around it.
- A number: digits only, no commas, no units ($, %, km) unless the question asks for them; if it asks for
  e.g. USD with two decimals, give exactly that (e.g. 1234.50).
- A string: as few words as possible, no articles, no abbreviations unless asked. Asked for a first name or
  surname only? Give only that.
- A list: comma separated (", "), in the order or sorting the question asks for.
- Follow any format the question gives exactly (e.g. 'IOC code', 'algebraic notation', 'alphabetical order')."""


def build_agent() -> CodeAgent:
    model = LiteLLMModel(
        model_id=f"gemini/{GEMINI_MODELS[0]}",
        api_key=os.environ["GEMINI_API_KEY"],
        num_retries=1,
        timeout=90,  # some preview models hang for minutes; skip to the next one instead
        reasoning_effort="low",  # default thinking spends thousands of tokens and minutes per step
        fallbacks=[f"gemini/{m}" for m in GEMINI_MODELS[1:]],  # litellm moves on when a model hits its quota
    )
    return CodeAgent(
        tools=[DuckDuckGoSearchTool(), visit_webpage, WikipediaSearchTool(user_agent=USER_AGENT), ask_about_file, ask_about_youtube, youtube_transcript],
        model=model,
        instructions=INSTRUCTIONS,
        additional_authorized_imports=["pandas", "openpyxl", "xlrd", "numpy", "math", "datetime", "re", "json", "itertools", "collections"],
        code_block_tags="markdown",  # Gemini often skips the closing </code> tag, so the step never stops
        max_steps=12,
        verbosity_level=1,
    )


def clean_answer(answer) -> str:
    text = str(answer).strip()
    text = re.sub(r"^(final answer|answer)\s*:\s*", "", text, flags=re.I)
    text = text.strip().strip('"').strip("'").strip()
    return text[:-1] if text.endswith(".") and not re.search(r"\d\.\d+$", text) else text


if __name__ == "__main__":
    assert clean_answer("FINAL ANSWER: Paris.") == "Paris"
    assert clean_answer(' "Rd5" ') == "Rd5"
    assert clean_answer(3.14) == "3.14"
    assert clean_answer(89706.00) == "89706.0"
    print("clean_answer ok")