import os

from smolagents import OpenAIServerModel, ToolCallingAgent

from tools import (
    AudioTranscriberTool,
    txt_reader,
    pdf_reader,
    excel_reader,
    math_calculator,
    search_engine,
    image_analyzer,
    YoutubeTranscriberTool,
    YoutubeVideoDescriptorTool,
)
from utils import OPENAI_MODEL_ID

MAIN_PROMPT = """
You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template:

[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.

Each question will contain the question, the task ID, and if available, a file name.

Key points:
- Do not add any prefix or suffix to YOUR FINAL ANSWER.
- Answer shortly and concisely, following the template strictly.
- Again, YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.

Example:
[Final Answer: 42]
Good example: "42"
Bad example: "The answer is 42"
"""


class MainAgent:
    """Main agent that orchestrates file handling and web search tasks."""

    def __init__(self):
        model = OpenAIServerModel(model_id=OPENAI_MODEL_ID, api_key=os.getenv("OPENAI_API_KEY"))
        self.agent = ToolCallingAgent(
            name="main_agent",
            description="An agent that can search the web, visit webpages, perform calculations, and handle files.",
            tools=[
                search_engine,
                math_calculator,
                AudioTranscriberTool(),
                txt_reader,
                pdf_reader,
                excel_reader,
                image_analyzer,
                YoutubeTranscriberTool(),
                YoutubeVideoDescriptorTool(),
            ],
            max_steps=5,
            model=model,
        )

    def run(self, question: str, task_id: str, file_name: str | None) -> str:
        """Run the agent with the provided question, task ID, and optional file name."""
        query = f"""q: {question}
task_id: {task_id}
{'file_name: ' + file_name if file_name else ''}
"""
        return self.agent.run(query)