import agent_tools
from smolagents import CodeAgent, AzureOpenAIServerModel
import app_tokens

model = AzureOpenAIServerModel(
    model_id = app_tokens.AZURE_OPENAI_MODEL,
    azure_endpoint = app_tokens.AZURE_OPENAI_ENDPOINT,
    api_key = app_tokens.AZURE_OPENAI_API_KEY,
    api_version = app_tokens.OPENAI_API_VERSION    
)

class BasicAgent:
    def __init__(self):
        self.web_agent = CodeAgent(
            model=model,
            tools=[#agent_tools.tavily_search, 
                   #agent_tools.duck_search, 
                   agent_tools.visit_page, 
                   agent_tools.wiki_search, 
                   agent_tools.google_search, 
                   agent_tools.final_answer
                   ],
            max_steps=8,
            name="web_agent",
            description="Runs web searches for you. Can use Google, DuckDuckGo and Wikipedia for search.",
            add_base_tools = True
        )

        self.audio_agent = CodeAgent(
            model=model,
            tools=[agent_tools.speech_to_text_tool, agent_tools.final_answer],
            max_steps=4,
            name="audio_agent",
            description="This agent can help you with converting audio to text.",
            add_base_tools = True
        )

        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="This agent can help you with executing and validating python code.",
            add_base_tools = True
        )

        self.visual_agent = CodeAgent(
            model=model,
            tools=[agent_tools.visual_qa_tool, agent_tools.final_answer],
            max_steps=4,
            name="visual_qa_agent",
            description="This agent can help you with answering questions about pictures.",
            add_base_tools = True
        )

        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,
            add_base_tools = True
        )

    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