import agent_tools
from smolagents import CodeAgent, AzureOpenAIServerModel
import my_tokens

# Instancia del modelo
model = AzureOpenAIServerModel(
    model_id=my_tokens.AZURE_OPENAI_MODEL,
    azure_endpoint=my_tokens.AZURE_OPENAI_ENDPOINT,
    api_key=my_tokens.AZURE_OPENAI_API_KEY,
    api_version=my_tokens.OPENAI_API_VERSION    
)

class CustomAgent:
    def __init__(self):
        # Definir agentes de trabajo con sus herramientas
        self.web_agent = CodeAgent(
            model=model,
            tools=[agent_tools.google_search, agent_tools.wiki_search, agent_tools.visit_page, agent_tools.final_answer],
            max_steps=8,
            name="web_agent",
            description="Este agente maneja búsquedas web."
        )

        self.audio_agent = CodeAgent(
            model=model,
            tools=[agent_tools.speech_to_text_tool, agent_tools.final_answer],
            max_steps=4,
            name="audio_agent",
            description="Este agente convierte audio a texto."
        )

        self.py_agent = CodeAgent(
            model=model,
            tools=[agent_tools.do_python, agent_tools.final_answer],
            additional_authorized_imports=["json", "pandas", "numpy", "regex"],
            max_steps=8,
            name="python_code_agent",
            description="Este agente ejecuta y valida código Python."
        )

        self.visual_agent = CodeAgent(
            model=model,
            tools=[agent_tools.visual_qa_tool, agent_tools.final_answer],
            max_steps=4,
            name="visual_qa_agent",
            description="Este agente responde preguntas sobre imágenes."
        )

        self.manager_agent = CodeAgent(
            model=model,
            tools=[],
            managed_agents=[self.web_agent, self.audio_agent, self.py_agent, self.visual_agent],
            planning_interval=8,
            verbosity_level=2,
            max_steps=12,
        )

    def forward(self, question: str, attachment: str = None) -> str:
        if attachment:
            result = self.manager_agent.run(question, additional_args={"attachment": attachment})
        else:
            result = self.manager_agent.run(question)
        return result