import os

from smolagents import CodeAgent, LiteLLMModel, WikipediaSearchTool

from prompt import SYSTEM_PROMPT
from tools import download_gaia_file, read_file_content


class BasicAgent:
    def __init__(self):
        self.model = LiteLLMModel(
            model_id="openrouter/anthropic/claude-sonnet-4",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            temperature=0.0,
            max_tokens=4096,
        )

        # add_base_tools=True gives us python_interpreter, web_search
        # (DuckDuckGo) and visit_webpage for free, so we only add what's
        # missing: Wikipedia search plus our own GAIA file tools.
        self.agent = CodeAgent(
            model=self.model,
            tools=[
                WikipediaSearchTool(user_agent="GAIA-BasicAgent (contact: example@example.com)"),
                download_gaia_file,
                read_file_content,
            ],
            add_base_tools=True,
            max_steps=10,
            instructions=SYSTEM_PROMPT,
            additional_authorized_imports=[
                "pandas",
                "numpy",
                "requests",
                "json",
                "re",
                "math",
                "statistics",
                "datetime",
                "collections",
                "itertools",
                "PIL",
            ],
        )

    def __call__(self, question: str, task_id: str | None = None) -> str:
        prompt = question
        if task_id:
            prompt = (
                f"task_id: {task_id}\n"
                f"Question: {question}\n\n"
                "If this question mentions or implies an attached file, call "
                "download_gaia_file with this exact task_id first, then call "
                "read_file_content on the path it returns before answering."
            )
        try:
            result = self.agent.run(prompt)
        except Exception as e:
            return f"AGENT ERROR: {e}"
        return str(result).strip()