import os
from smolagents import CodeAgent, DuckDuckGoSearchTool, tool, InferenceClientModel, OpenAIServerModel, SpeechToTextTool, WikipediaSearchTool
from dotenv import load_dotenv
load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

@tool
def add(a: int, b: int) -> int:
    """
    Adds two numbers together.
    Args:
        a (int): The first number.
        b (int): The second number.
    Returns:
        int: The result of a + b.
    """
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """
    Subtracts the second number from the first.
    Args:
        a (int): The number to subtract from.
        b (int): The number to subtract.
    Returns:
        int: The result of a - b.
    """
    return a - b

@tool
def multiply(a: int, b: int) -> int:
    """
    Multiplies two numbers.
    Args:
        a (int): The first number.
        b (int): The second number.
    Returns:
        int: The product of a and b.
    """
    return a * b

@tool
def divide(a: int, b: int) -> float:
    """
    Divides the first number by the second.
    Args:
        a (int): The numerator.
        b (int): The denominator.
    Returns:
        float: The result of a / b.
    Raises:
        ZeroDivisionError: If b is zero.
    """
    return a / b

prompt = "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."

class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")
        llm = OpenAIServerModel(
        model_id="gpt-4o",
        api_base="https://api.openai.com/v1",
        api_key=os.environ["OPENAI_API_KEY"],
        )
        
        self.agent = CodeAgent(
            tools=[add,subtract,multiply,divide,DuckDuckGoSearchTool(),WikipediaSearchTool(),SpeechToTextTool()],
            model=llm,
        )
    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        # Prepend the prompt to the question with a clear separator
        full_input = prompt + "\n\nQuestion: " + question
        fixed_answer = self.agent.run(full_input)
        print(f"Agent returning fixed answer: {fixed_answer}")
        return fixed_answer