"""
agent.py ­– Gemini-smolagents baseline using google-genai SDK
------------------------------------------------------------
Environment variables
---------------------
GOOGLE_API_KEY            API key from Google AI Studio
Optional:
    GAIA_API_URL          GAIA evaluation endpoint (default: official URL)

This file defines:
    • GeminiModel  – wraps google-genai for smolagents
    • gaia_file_reader  – custom tool to fetch <file:xyz> attachments
    • GeminiAgent  – CodeAgent with Python / Search / File tools + Gemini model
"""

import os
import re
import base64
import mimetypes
import requests
import google.genai as genai
from google.genai import types as gtypes
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    PythonInterpreterTool,
    tool,
)

# --------------------------------------------------------------------------- #
# Constants & helpers
# --------------------------------------------------------------------------- #
DEFAULT_API_URL = os.getenv(
    "GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space"
)
FILE_TAG = re.compile(r"<file:([^>]+)>")


def _download_file(file_id: str) -> bytes:
    """Download the attachment for a GAIA task."""
    url = f"{DEFAULT_API_URL}/files/{file_id}"
    resp = requests.get(url, timeout=30)
    resp.raise_for_status()
    return resp.content


# --------------------------------------------------------------------------- #
# Model wrapper
# --------------------------------------------------------------------------- #
class GeminiModel:
    """
    Thin adapter around google-genai Client for smolagents.
    """

    def __init__(
        self,
        model_name: str = "gemini-2.0-flash",
        temperature: float = 0.1,
        max_tokens: int = 128,
    ):
        api_key = os.getenv("GOOGLE_API_KEY")
        if not api_key:
            raise EnvironmentError("GOOGLE_API_KEY is not set.")
        self.client = genai.Client(api_key=api_key)
        self.model_name = model_name
        self.temperature = temperature
        self.max_tokens = max_tokens

    # ---------- main generation helpers ---------- #
    def call(self, prompt: str, **kwargs) -> str:
        """Text-only helper used by __call__."""
        resp = self.client.models.generate_content(
            model=self.model_name,
            contents=prompt,
            config=gtypes.GenerateContentConfig(
                temperature=self.temperature,
                max_output_tokens=self.max_tokens,
            ),
        )
        return resp.text.strip()

    def call_messages(self, messages, **kwargs) -> str:
        sys_msg, user_msg = messages
        contents = (
            [sys_msg["content"], *user_msg["content"]]
            if isinstance(user_msg["content"], list)
            else f"{sys_msg['content']}\n\n{user_msg['content']}"
        )
        resp = self.client.models.generate_content(
            model=self.model_name,
            contents=contents,
            config=gtypes.GenerateContentConfig(
                temperature=self.temperature,
                max_output_tokens=self.max_tokens,
            ),
        )
        return resp.text.strip()

    # ---------- make the instance itself callable ---------- #
    def __call__(self, prompt: str, **kwargs) -> str:  # <-- NEW
        return self.call(prompt, **kwargs)


# --------------------------------------------------------------------------- #
# Custom tool: fetch GAIA attachments
# --------------------------------------------------------------------------- #
@tool
def gaia_file_reader(file_id: str) -> str:
    """
    Download a GAIA attachment and return its contents.

    Args:
        file_id: The identifier that appears inside a <file:...> placeholder
                 in the GAIA question prompt.

    Returns:
        A base-64 string for binary files (images, PDF, etc.) or UTF-8 text for
        plain-text files.
    """
    try:
        raw = _download_file(file_id)
        mime = mimetypes.guess_type(file_id)[0] or "application/octet-stream"
        if mime.startswith("text") or mime in ("application/json",):
            return raw.decode(errors="ignore")
        return base64.b64encode(raw).decode()
    except Exception as exc:
        return f"ERROR downloading {file_id}: {exc}"


# --------------------------------------------------------------------------- #
# Final agent class
# --------------------------------------------------------------------------- #
class GeminiAgent:
    def __init__(self):
        self.system_prompt = (
            "You are a concise, highly accurate assistant. "
            "Unless explicitly required, reply with ONE short sentence. "
            "Use the provided tools if needed. "
            "All answers are graded by exact string match."
        )

        model = GeminiModel()
        tools = [
            PythonInterpreterTool(),
            DuckDuckGoSearchTool(),
            gaia_file_reader,
        ]

        # ✨  system_prompt removed – newest smolagents doesn't take it
        self.agent = CodeAgent(
            model=model,
            tools=tools,
            # any other kwargs (executor_type, additional_authorized_imports…)
            verbosity_level=0,
        )
        print("✅ GeminiAgent ready.")

    def __call__(self, question: str) -> str:
        file_ids = FILE_TAG.findall(question)

        # -------- multimodal branch -------- #
        if file_ids:
            parts: list[gtypes.Part] = []
            text_part = FILE_TAG.sub("", question).strip()
            if text_part:
                parts.append(gtypes.Part.from_text(text_part))
            for fid in file_ids:
                try:
                    img_bytes = _download_file(fid)
                    mime = mimetypes.guess_type(fid)[0] or "image/png"
                    parts.append(
                        gtypes.Part.from_bytes(data=img_bytes, mime_type=mime)
                    )
                except Exception as exc:
                    parts.append(
                        gtypes.Part.from_text(f"[FILE {fid} ERROR: {exc}]")
                    )
            messages = [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": parts},
            ]
            answer = self.agent.model.call_messages(messages)

        # -------- text-only branch -------- #
        else:
            # prepend system prompt to the user question
            full_prompt = f"{self.system_prompt}\n\n{question}"
            answer = self.agent(full_prompt)

        return answer.rstrip(" .\n\r\t")