from math import sqrt
import logging
from llama_index.core.tools import FunctionTool

def add(a: float, b: float, context: str) -> str:
    """Add two numbers.
    Args:
        a (float): Addend 1
        b (float): Addend 2
        context (str): Context for the addition, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a + b)} for {context}"


def sub(a: float, b: float, context: str) -> str:
    """Subtract two numbers.
    Args:
        a (float): Minuend
        b (float): Subtrahend
        context (str): Context for the subtraction, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a - b)} for {context}"


def multiply(a: float, b: float, context: str) -> str:
    """Multiply two numbers.
    Args:
        a (float): Multiplicand
        b (float): Multiplier
        context (str): Context for the multiplication, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a * b)} for {context}"


def divide(a: float, b: float, context: str) -> str:
    """Divide two numbers.
    Args:
        a (float): Dividend
        b (float): Divisor
        context (str): Context for the division, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included, or an error message if division by zero occurs.
    """
    if b == 0:
        return f"Error: Division by zero for {context}"
    return f"Result {str(a / b)} for {context}"


def floordiv(a: float, b: float, context: str) -> str:
    """Floor divide two numbers.
    Args:
        a (float): Dividend
        b (float): Divisor
        context (str): Context for the floor division, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included, or an error message if division by zero occurs.
    """
    if b == 0:
        return f"Error: Division by zero for {context}"
    return f"Result {str(a // b)} for {context}"


def mod(a: float, b: float, context: str) -> str:
    """Get the modulus of two numbers.
    Args:
        a (float): Dividend
        b (float): Divisor
        context (str): Context for the modulus operation, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a % b)} for {context}"


def sqrt(a: float, context: str) -> str:
    """Get the square root of a number.
    Args:
        a (float): Number to find the square root of
        context (str): Context for the square root operation, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a**0.5)} for {context}"


def equals(a: float, b: float, context: str) -> str:
    """Check if two numbers are equal.
    Args:
        a (float): First number
        b (float): Second number
        context (str): Context for the equality check, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a == b)} for {context}"


def less_than(a: float, b: float, context: str) -> str:
    """Check if the first number is less than the second number.
    Args:
        a (float): First number
        b (float): Second number
        context (str): Context for the less than check, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a < b)} for {context}"


def more_than(a: float, b: float, context: str) -> str:
    """Check if the first number is more than the second number.
    Args:
        a (float): First number
        b (float): Second number
        context (str): Context for the more than check, this should be used to keep track of the operations.
    Returns:
        str: The result with the context included.
    """
    return f"Result {str(a > b)} for {context}"

def solve_math_problem(problem: str) -> str:
    """
    Solve mathematical problems using GPT-4.
    Args:
        problem: The mathematical problem to solve
    Returns:
        str: Solution
    """
    
    try:
        from openai import OpenAI
        client = OpenAI()
        
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": "You are a mathematics expert. Solve problems step by step with clear explanations."
                },
                {
                    "role": "user",
                    "content": f"Please solve this mathematical problem step by step:\n\n{problem}"
                }
            ]
        )
        
        return response.choices[0].message.content
        
    except Exception as e:
        return f"Math problem solving failed: {str(e)}"


def get_math_tools():
    """
    Get the math tools.
    """
    return [
        FunctionTool.from_defaults(fn=sum),
        FunctionTool.from_defaults(fn=sub),
        FunctionTool.from_defaults(fn=add),
        FunctionTool.from_defaults(fn=sqrt),
        FunctionTool.from_defaults(fn=floordiv),
        FunctionTool.from_defaults(fn=mod),
        FunctionTool.from_defaults(fn=multiply),
        FunctionTool.from_defaults(fn=equals),
        FunctionTool.from_defaults(fn=less_than),
        FunctionTool.from_defaults(fn=more_than),
    ]