from dotenv import load_dotenv
from smolagents import CodeAgent
from smolagents import OpenAIServerModel
from tool import fetch_webpage, read_file_tool, get_youtube_transcript, transcribe_audio

from smolagents import VisitWebpageTool, WikipediaSearchTool, PythonInterpreterTool, DuckDuckGoSearchTool, WebSearchTool, SpeechToTextTool

from prompt import gaia_prompt

load_dotenv()

openai_nano_model = OpenAIServerModel(
    # model_id="gpt-4.1-nano-2025-04-14",
    model_id="gpt-4.1-mini-2025-04-14",
    # model_id="o3-mini-2025-01-31",
)

gaia_agent = CodeAgent(
    model=openai_nano_model,
    tools=[fetch_webpage, DuckDuckGoSearchTool(), PythonInterpreterTool(), read_file_tool, get_youtube_transcript, transcribe_audio], # WikipediaSearchTool(), VisitWebpageTool(max_output_length=60000)
    max_steps=7,
    additional_authorized_imports=["requests", "bs4", "pandas", "numpy", "markdownify"]
)

class GAIA_Agent:
    def __init__(self):
        self.system_prompt = gaia_prompt
        self.agent = gaia_agent

    def __call__(self, question: str) -> str:

        try:

            full_context = self.system_prompt + "\nTHE QUESTION:\n" + question

            final_answer = self.agent.run(full_context)
            print(f"final answer returned by agent ----> {final_answer}")
            return final_answer
        except Exception as e:
            error = f"An error occurred while processing the question: {e}"
            print(error)
            return error

# build context + append instructions and all

# clean answer function

if __name__ == "__main__":
    answer = gaia_agent.run(
        """
        You are a general AI assistant. I will ask you a question. You can 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. Remember: GAIA requires exact answer matching. Just provide the factual answer.

        How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."""
    )
    pass