import os

from openai import OpenAI

from smolagents import (
    CodeAgent, 
    OpenAIServerModel, 
    DuckDuckGoSearchTool, 
    WikipediaSearchTool
)

MODEL_ID = "o4-mini"

def run_agent(level, question, file_name):
    if level == 1:
        max_steps = 5
    elif level == 2:
        max_steps = 10
    elif level == 3:
        max_steps = 20
   
    agent = CodeAgent(
            add_base_tools = True,
            additional_authorized_imports = ["pandas", "numpy"],
            executor_type = "e2b", # Remote sandbox: https://e2b.dev/dashboard/
            executor_kwargs = {"api_key": os.getenv("E2B_API_KEY")},
            max_steps = max_steps,
            model = OpenAIServerModel(model_id = MODEL_ID),
            #planning_interval=3,
            tools = [#DuckDuckGoSearchTool(), 
                     WikipediaSearchTool()],
            verbosity_level = 1
        )

    agent.logger.console.width = 66
    
    answer = agent.run(question)

    return get_final_answer(question, answer)

def get_final_answer(question, answer):
    client = OpenAI()

    prompt_template = """
You are a GAIA benchmark assistant, who is given a question and an answer. Read the question carefully. You must give the final answer requested. Do not include explanations, steps, reasoning, or additional text.
**Question:** """ + str(question) + """
**Answer:** """ + str(answer) + """
**Final answer:**:

"""

    completion = client.chat.completions.create(
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt_template}]}],
        model = MODEL_ID
    )

    final_answer = completion.choices[0].message.content

    return final_answer