from abc import abstractmethod, ABC
from smolagents import CodeAgent, Tool
from agent.constants import ADDITIONAL_AUTHORIZED_IMPORT
from tools.tools_utils import ToolsUtils


class BaseAgent(ABC):
    def __init__(self, model_name: str, tools: list[Tool] | None = None, planning_interval: int = 3, max_steps: int = 12, use_all_custom_tools: bool = True):
        self.model_name: str = model_name
        self.planning_interval = planning_interval
        self.max_steps = max_steps
        self.use_all_custom_tools = use_all_custom_tools
        self.tools: list[Tool] = self.init_tools(tools)
        self.agent = self.init_agent()

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        fixed_answer = self.agent.run(question)
        print(f"Agent returning fixed answer: {fixed_answer}")
        return fixed_answer

    @abstractmethod
    def get_model(self):
        pass

    def init_tools(self, tools: list[Tool] | None = None):
        if tools is None:
            tools = []
        if self.use_all_custom_tools:
            tools = ToolsUtils.get_default_tools()
        return tools

    def add_tool(self, tool: Tool):
        self.tools.append(tool)

    def init_agent(self):
        agent = CodeAgent(
            model=self.get_model(),
            tools=[t for t in self.tools],
            add_base_tools=True,
            verbosity_level=2,
            additional_authorized_imports=ADDITIONAL_AUTHORIZED_IMPORT,
            planning_interval=self.planning_interval,
            max_steps=self.max_steps
        )
        agent.system_prompt += ("\nWhen answering, provide ONLY the precise answer requested. "
                                "Do not include explanations, steps, reasoning, or additional text. "
                                "Be direct and specific. GAIA benchmark requires exact matching answers.")
        # print(agent.system_prompt)
        print("Agent initialized.")
        return agent