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

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="o3-mini-2025-01-31",
)

gaia_agent = CodeAgent(
    model=openai_nano_model,
    tools=[fetch_webpage, DuckDuckGoSearchTool(), PythonInterpreterTool(), read_file_tool, get_youtube_transcript], # WikipediaSearchTool(), VisitWebpageTool(max_output_length=60000)
    max_steps=5,
    verbosity_level=2,
    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)
            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__":
    pass
    # gaia_agent.run("What is the weather in Mumbai?")
#     answer = gaia_agent.run(
# f"""
# 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.
# """
#     )

    # print(gaia_prompt)
#     answer = gaia_agent.run("""
#         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.
#         You can search for results and then visit a webpage to get more information. Break down the problem into smaller sub-problems and solve them one by one.
#         Think like a human.

#         What is the final numeric output from the attached Python code?
                            
#         ----

# """)
#     print(f"this is the final answer the gaia agent gave ---> {answer}")