import math
from typing import Optional, Tuple, Literal
from smolagents import tool



@tool
def real_number_calculator(
    a: float, b: float, operation: Literal["add", "subtract", "multiply", "divide"]
) -> float:
    """
    Perform basic arithmetic operations on two real numbers.

    This function acts as a simple calculator that can add, subtract, 
    multiply, or divide two floating-point numbers.

    Args:
        a (float): The first number.
        b (float): The second number.
        operation (Literal["add", "subtract", "multiply", "divide"]): 
            The arithmetic operation to perform.

    Returns:
        float: The result of the arithmetic operation.

    Raises:
        ValueError: If the operation is invalid or if division by zero is attempted.

    Example:
        >>> real_number_calculator(10, 5, "add")
        15.0
        >>> real_number_calculator(10, 5, "divide")
        2.0
    """

    if operation == "add":
        return a + b
    elif operation == "subtract":
        return a - b
    elif operation == "multiply":
        return a * b
    elif operation == "divide":
        if b == 0:
            raise ValueError("Division by zero is not allowed.")
        return a / b
    else:
        raise ValueError(f"Invalid operation: {operation}")




class TestAgent:
    def __init__(self):
        
        # import code agent and basic tool from smolagent
        from smolagents import CodeAgent, OpenAIServerModel, DuckDuckGoSearchTool, FinalAnswerTool, VisitWebpageTool, MCPClient

        # import additional tool from langchain @ https://docs.langchain.com/oss/python/integrations/tools
        #from langchain_community.agent_toolkits import load_tools
        from langchain_community.agent_toolkits.load_tools import load_tools

        from smolagents import Tool
        wikipedia_tool = Tool.from_langchain(load_tools(["wikipedia"])[0])
        wikipedia_tool.top_k_results=3

        # import tools from MCP servers @ https://github.com/mcp
        #from mcp import StdioServerParameters
        #server_parameters = StdioServerParameters(command="uvx",
        #                                          args=["--quiet", "youtubeqa@0.2.1"],
        #                                          env={"UV_PYTHON": "3.12", **os.environ},
        #                                         )
        #youtube_tools = MCPServerTool(server_params=server_parameters)

        model = OpenAIServerModel(model_id="gpt-4o")
        #model = InferenceClientModel("Qwen/Qwen2.5-Coder-32B-Instruct")
        # Instantiate the agent
        self.agent = CodeAgent(
            tools=[real_number_calculator,        # homemade tool
                   DuckDuckGoSearchTool(),          # basic tools from smolagent
                   VisitWebpageTool(),
                   wikipedia_tool,                  # tool from langchain with extra parmaeters
                   #youtube_tools,                   # tool from MCP server
                   FinalAnswerTool()],
            additional_authorized_imports=["pandas","markdownify","requests"],
            model=model,
            max_steps=3,
            verbosity_level=2,
            use_structured_outputs_internally=True
        )

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

        print(f"Agent received question (first 50 chars): {question[:50]}...")
        answer = self.agent.run(question)
        print(f"Agent returning his answer: {answer}")
        return answer