from smolagents import InferenceClientModel, CodeAgent, GoogleSearchTool, VisitWebpageTool, WikipediaSearchTool

from tool.SupadataTool import SupadataTool
from tool.UnknownTextTool import UnknownTextTool, BasicMathTool


class GaiaAgent:
    def __init__(self):
        manager_model = InferenceClientModel(
            model_id='Qwen/Qwen3-30B-A3B-Thinking-2507',
            provider='nebius',
            max_tokens=8096,
        )
        agent_model = InferenceClientModel(
            "mistralai/Devstral-Small-2505",
            provider="nebius",
            max_tokens=8096
        )

        self.web_agent = CodeAgent(
            model=agent_model,
            tools=[
                GoogleSearchTool(provider="serper"),
                VisitWebpageTool()
            ],
            name="web_agent",
            description="Browses the web to find information",
            verbosity_level=0,
            max_steps=10,
        )
        self.wikipedia_agent = CodeAgent(
            model=agent_model,
            tools=[WikipediaSearchTool()],
            name="wikipedia_agent",
            description="Searches Wikipedia and returns a summary or full text of the given topic, along with the page URL. If you need to get any encyclopedia information, use this tool.",
            max_steps=5,
        )
        self.text_agent = CodeAgent(
            model=agent_model,
            tools=[UnknownTextTool(), BasicMathTool()],
            name="text_agent",
            description="This agent is perfect for solving specific tasks: * decoding unknown text (if you have unknown text, use this tool); * math and algorithmic problems add, subtract, multiply, divide, etc.",
        )
        self.youtube_transcript_agent = CodeAgent(
            model=agent_model,
            tools=[SupadataTool()],
            name="youtube_transcript_agent",
            description="This agent is perfect for generating transcripts of YouTube videos. If you need to know what the video is about, use this tool.",
        )
        self.manager_agent = CodeAgent(
            tools=[],
            model=manager_model,
            managed_agents=[self.web_agent, self.wikipedia_agent, self.text_agent, self.youtube_transcript_agent],
            additional_authorized_imports=[
            ],
            planning_interval=5,
            verbosity_level=2,
            max_steps=15,
        )

        print("GaiaAgent initialized.")


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