import re
from smolagents import AgentMemory, CodeAgent, InferenceClientModel, FinalAnswerTool, WebSearchTool
from collections.abc import Callable

from smolagents.default_tools import VisitWebpageTool, WikipediaSearchTool

class MultiAgent:
    def __init__(self):

        # Instructions prompt
        self.instructions = """finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
        YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
        If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
        If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""

        # WEB AGENT
        self.web_agent = CodeAgent(
            model=InferenceClientModel("Qwen/Qwen2.5-Coder-32B-Instruct", provider="nebius", max_tokens=8096),
            tools=[WikipediaSearchTool(), WebSearchTool(), VisitWebpageTool()],
            name="web_agent",
            description="A code agent that can search the web and visit webpages",
            additional_authorized_imports=[ "markdownify" , "requests" ],
            max_steps=5
        )

        # MANAGER AGENT
        self.manager_agent = CodeAgent(
            model=InferenceClientModel("deepseek-ai/DeepSeek-R1", provider="nebius", max_tokens=8096),
            tools=[FinalAnswerTool()],
            managed_agents=[self.web_agent],
            instructions=self.instructions,
            additional_authorized_imports=[
                "pandas",
                "numpy",
                "chess"
            ],
            planning_interval=5,
            verbosity_level=2,
            max_steps=5,
        )

        print("Multi Agent 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 answer: {answer}")
        return answer

    def final_answer_checks(self) -> list[Callable] :
        return [ self.check_func ]

    def check_func(self, answer: str, memory: AgentMemory) -> bool:
        check = bool(re.match(r'^(\d+(\.\d+)?|\w+(\s+\w+){0,4}|(\d+(\.\d+)?|"[^"]*"|\w+)(\s*,\s*(\d+(\.\d+)?|"[^"]*"|\w+))+)$', answer))
        print(f"FINAL ANSWER CHECK is {check}")
        return check