from smolagents import (
    ToolCallingAgent,
    CodeAgent,
    DuckDuckGoSearchTool,
    PythonInterpreterTool,
    FinalAnswerTool,
    VisitWebpageTool,
    PromptTemplates,
    PlanningPromptTemplate,
    ManagedAgentPromptTemplate,
    FinalAnswerPromptTemplate,
    WikipediaSearchTool,
)

from tool import addition_tool
from src.config.model import openai_model
import re


gaia_prompt = "You are a general AI assistant. I will ask you a question. Report your thoughts, and 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 then reply with digit and not in words, e.g: if Five, then reply with 5, if Three, then reply with 3, 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."
# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
class BaseAgent:
    def __init__(self):
        print("BaseAgent initialized.")
        self.agent = ToolCallingAgent(
            model=openai_model,
            tools=[
                addition_tool,
                WikipediaSearchTool(),
                PythonInterpreterTool(),
                DuckDuckGoSearchTool(),
                VisitWebpageTool(),
                FinalAnswerTool(),
            ],
            add_base_tools=True,
            max_steps=4,
            verbosity_level=2,
            prompt_templates=PromptTemplates(
                system_prompt=gaia_prompt,
                planning=PlanningPromptTemplate(
                    initial_plan="",
                    update_plan_pre_messages="",
                    update_plan_post_messages="",
                ),
                managed_agent=ManagedAgentPromptTemplate(
                    task="",
                    report="",
                ),
                final_answer=FinalAnswerPromptTemplate(
                    pre_messages=gaia_prompt, post_messages=gaia_prompt+"The final answer should be as few words or numbers only. Exclude all the previuos context and only return the final answer. In case words separated by commas, use the following format: 'word1, word2, word3' without any additional text or explanation.",
                ),
            ),
        )

    def run(self, question: str, additional_args: dict = None) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        response = self.agent.run(question, additional_args=additional_args)
        response = re.sub(r"(?s)^.*?FINAL ANSWER:", "FINAL ANSWER:", response)
        response = response.replace("FINAL ANSWER:", "").strip()
        print(f"Agent returning response: {response}")
        return response