import json
import math

from agentica.tools.base import Tool
from agentica.utils.log import logger


class CalculatorTool(Tool):
    def __init__(
            self,
            add: bool = True,
            subtract: bool = True,
            multiply: bool = True,
            divide: bool = True,
            exponentiate: bool = False,
            factorial: bool = False,
            is_prime: bool = False,
            square_root: bool = False,
            enable_all: bool = False,
    ):
        super().__init__(name="calculator_tool")

        # Register functions in the toolkit
        if add or enable_all:
            self.register(self.add)
        if subtract or enable_all:
            self.register(self.subtract)
        if multiply or enable_all:
            self.register(self.multiply)
        if divide or enable_all:
            self.register(self.divide)
        if exponentiate or enable_all:
            self.register(self.exponentiate)
        if factorial or enable_all:
            self.register(self.factorial)
        if is_prime or enable_all:
            self.register(self.is_prime)
        if square_root or enable_all:
            self.register(self.square_root)

    def add(self, a: float, b: float) -> str:
        """Add two numbers and return the result.

        Args:
            a (float): First number.
            b (float): Second number.

        Returns:
            str: JSON string of the result.

        Example:
            from agentica.tools.calculator_tool import CalculatorTool
            m = CalculatorTool()
            r = m.add(1, 2)
            print(r)
        """
        result = a + b
        logger.info(f"Adding {a} and {b} to get {result}")
        return json.dumps({"operation": "addition", "result": result})

    def subtract(self, a: float, b: float) -> str:
        """Subtract second number from first and return the result.

        Args:
            a (float): First number.
            b (float): Second number.

        Returns:
            str: JSON string of the result.
        """
        result = a - b
        logger.info(f"Subtracting {b} from {a} to get {result}")
        return json.dumps({"operation": "subtraction", "result": result})

    def multiply(self, a: float, b: float) -> str:
        """Multiply two numbers and return the result.

        Args:
            a (float): First number.
            b (float): Second number.

        Returns:
            str: JSON string of the result.
        """
        result = a * b
        logger.info(f"Multiplying {a} and {b} to get {result}")
        return json.dumps({"operation": "multiplication", "result": result})

    def divide(self, a: float, b: float) -> str:
        """Divide first number by second and return the result.

        Args:
            a (float): Numerator.
            b (float): Denominator.

        Returns:
            str: JSON string of the result.
        """
        if b == 0:
            logger.error("Attempt to divide by zero")
            return json.dumps({"operation": "division", "error": "Division by zero is undefined"})
        try:
            result = a / b
        except Exception as e:
            return json.dumps({"operation": "division", "error": e, "result": "Error"})
        logger.info(f"Dividing {a} by {b} to get {result}")
        return json.dumps({"operation": "division", "result": result})

    def exponentiate(self, a: float, b: float) -> str:
        """Raise first number to the power of the second number and return the result.

        Args:
            a (float): Base.
            b (float): Exponent.

        Returns:
            str: JSON string of the result.
        """
        result = math.pow(a, b)
        logger.info(f"Raising {a} to the power of {b} to get {result}")
        return json.dumps({"operation": "exponentiation", "result": result})

    def factorial(self, n: int) -> str:
        """Calculate the factorial of a number and return the result.

        Args:
            n (int): Number to calculate the factorial of.

        Returns:
            str: JSON string of the result.
        """
        if n < 0:
            logger.error("Attempt to calculate factorial of a negative number")
            return json.dumps({"operation": "factorial", "error": "Factorial of a negative number is undefined"})
        result = math.factorial(n)
        logger.info(f"Calculating factorial of {n} to get {result}")
        return json.dumps({"operation": "factorial", "result": result})

    def is_prime(self, n: int) -> str:
        """Check if a number is prime and return the result.

        Args:
            n (int): Number to check if prime.

        Returns:
            str: JSON string of the result.
        """
        if n <= 1:
            return json.dumps({"operation": "prime_check", "result": False})
        for i in range(2, int(math.sqrt(n)) + 1):
            if n % i == 0:
                return json.dumps({"operation": "prime_check", "result": False})
        return json.dumps({"operation": "prime_check", "result": True})

    def square_root(self, n: float) -> str:
        """Calculate the square root of a number and return the result.

        Args:
            n (float): Number to calculate the square root of.

        Returns:
            str: JSON string of the result.
        """
        if n < 0:
            logger.error("Attempt to calculate square root of a negative number")
            return json.dumps({"operation": "square_root", "error": "Square root of a negative number is undefined"})

        result = math.sqrt(n)
        logger.info(f"Calculating square root of {n} to get {result}")
        return json.dumps({"operation": "square_root", "result": result})


if __name__ == '__main__':
    m = CalculatorTool()
    r = m.add(1, 2)
    print(r)
    r = m.subtract(1, 2)
    print(r)
    r = m.multiply(1, 2)
    print(r)
    r = m.divide(1, 2)
    print(r)
    r = m.exponentiate(1, 2)
    print(r)
    r = m.factorial(5)
    print(r)
    r = m.is_prime(5)
    print(r)
    r = m.square_root(5)
    print(r)
