from smolagents import (
    ToolCallingAgent,
    PromptTemplates,
    PlanningPromptTemplate,
    ManagedAgentPromptTemplate,
    FinalAnswerPromptTemplate,
)

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


# --- 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,
            ],
            verbosity_level=2,
            prompt_templates=PromptTemplates(
                system_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, 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.",
                # planning=PlanningPromptTemplate(
                #     initial_plan="",
                #     update_plan_pre_messages="",
                #     update_plan_post_messages="",
                # ),
                # managed_agent=ManagedAgentPromptTemplate(
                #     task="",
                #     report="",
                # ),
                # final_answer=FinalAnswerPromptTemplate(
                #     pre_messages="", post_messages=""
                # ),
            ),
        )

    def run(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        response = self.agent.run(question)
        print(f"Agent returning response: {response}")
        return response.replace("FINAL ANSWER:", "").strip()