from mcp.server.fastmcp import FastMCP
from typing import List, Optional
import json
from datetime import datetime
import math

# Initialize FastMCP server
mcp = FastMCP(
    "MathComputer",  # Name of the MCP server
    instructions="You are a mathematical computation assistant that can perform various mathematical operations and calculations.",
    host="0.0.0.0",
    port=8003,
)

@mcp.tool()
async def compute_operation(numbers: List[float], operation: str = "sum") -> str:
    """
    Perform mathematical operations on a list of numbers.
    
    Args:
        numbers (List[float]): List of numbers to operate on
        operation (str): Mathematical operation to perform
        
    Returns:
        str: Computation results
    """
    try:
        if not numbers:
            return "Error: Numbers list cannot be empty"
        
        if not all(isinstance(x, (int, float)) for x in numbers):
            return "Error: All values must be numerical"
        
        result = 0
        operation_details = {}
        
        if operation == "sum":
            result = sum(numbers)
            operation_details = {
                "method": "sum", 
                "count": len(numbers),
                "formula": "sum(numbers)"
            }
        elif operation == "average":
            result = sum(numbers) / len(numbers)
            operation_details = {
                "method": "average", 
                "count": len(numbers),
                "formula": "sum(numbers) / len(numbers)"
            }
        elif operation == "max":
            result = max(numbers)
            operation_details = {
                "method": "max", 
                "count": len(numbers),
                "formula": "max(numbers)"
            }
        elif operation == "min":
            result = min(numbers)
            operation_details = {
                "method": "min", 
                "count": len(numbers),
                "formula": "min(numbers)"
            }
        elif operation == "product":
            result = math.prod(numbers)
            operation_details = {
                "method": "product", 
                "count": len(numbers),
                "formula": "math.prod(numbers)"
            }
        elif operation == "geometric_mean":
            if any(x <= 0 for x in numbers):
                return "Error: Geometric mean requires all positive numbers"
            result = math.pow(math.prod(numbers), 1/len(numbers))
            operation_details = {
                "method": "geometric_mean", 
                "count": len(numbers),
                "formula": "pow(product, 1/n)"
            }
        elif operation == "harmonic_mean":
            if any(x == 0 for x in numbers):
                return "Error: Harmonic mean requires all non-zero numbers"
            result = len(numbers) / sum(1/x for x in numbers)
            operation_details = {
                "method": "harmonic_mean", 
                "count": len(numbers),
                "formula": "n / sum(1/x)"
            }
        else:
            return f"Error: Unsupported operation '{operation}'. Supported operations: sum, average, max, min, product, geometric_mean, harmonic_mean"
        
        computation_result = {
            "status": "success",
            "result": result,
            "operation": operation,
            "service_name": "MathComputer",
            "details": operation_details,
            "computation_timestamp": datetime.now().isoformat(),
            "input_numbers": numbers,
            "input_count": len(numbers)
        }
        
        return json.dumps(computation_result, indent=2, ensure_ascii=False)
        
    except Exception as e:
        return f"Error during computation: {str(e)}"

@mcp.tool()
async def get_supported_operations() -> str:
    """
    Get list of supported mathematical operations.
    
    Returns:
        str: Supported operations and descriptions
    """
    operations_info = {
        "service_name": "MathComputer",
        "supported_operations": [
            "sum", 
            "average", 
            "max", 
            "min", 
            "product", 
            "geometric_mean", 
            "harmonic_mean"
        ],
        "descriptions": {
            "sum": "Calculate the sum of all numbers",
            "average": "Calculate the arithmetic mean",
            "max": "Find the maximum value",
            "min": "Find the minimum value",
            "product": "Calculate the product of all numbers",
            "geometric_mean": "Calculate the geometric mean (requires positive numbers)",
            "harmonic_mean": "Calculate the harmonic mean (requires non-zero numbers)"
        },
        "description": "Mathematical computation service with advanced statistical functions"
    }
    
    return json.dumps(operations_info, indent=2, ensure_ascii=False)

@mcp.tool()
async def advanced_math_operations(operation: str, numbers: List[float], **kwargs) -> str:
    """
    Perform advanced mathematical operations.
    
    Args:
        operation (str): Advanced operation to perform
        numbers (List[float]): List of numbers
        **kwargs: Additional parameters for specific operations
        
    Returns:
        str: Advanced computation results
    """
    try:
        if not numbers:
            return "Error: Numbers list cannot be empty"
        
        if operation == "percentile":
            percentile = kwargs.get("percentile", 50)
            if not 0 <= percentile <= 100:
                return "Error: Percentile must be between 0 and 100"
            
            sorted_numbers = sorted(numbers)
            index = (percentile / 100) * (len(sorted_numbers) - 1)
            if index.is_integer():
                result = sorted_numbers[int(index)]
            else:
                lower = sorted_numbers[int(index)]
                upper = sorted_numbers[int(index) + 1]
                result = lower + (upper - lower) * (index - int(index))
            
            operation_details = {
                "method": "percentile",
                "percentile": percentile,
                "count": len(numbers),
                "formula": f"percentile_{percentile}(sorted_numbers)"
            }
            
        elif operation == "standard_deviation":
            if len(numbers) < 2:
                return "Error: Standard deviation requires at least 2 numbers"
            
            mean = sum(numbers) / len(numbers)
            variance = sum((x - mean) ** 2 for x in numbers) / (len(numbers) - 1)
            result = math.sqrt(variance)
            
            operation_details = {
                "method": "standard_deviation",
                "count": len(numbers),
                "formula": "sqrt(sum((x - mean)²) / (n-1))"
            }
            
        elif operation == "variance":
            if len(numbers) < 2:
                return "Error: Variance requires at least 2 numbers"
            
            mean = sum(numbers) / len(numbers)
            result = sum((x - mean) ** 2 for x in numbers) / (len(numbers) - 1)
            
            operation_details = {
                "method": "variance",
                "count": len(numbers),
                "formula": "sum((x - mean)²) / (n-1)"
            }
            
        else:
            return f"Error: Unsupported advanced operation '{operation}'. Supported: percentile, standard_deviation, variance"
        
        computation_result = {
            "status": "success",
            "result": result,
            "operation": operation,
            "service_name": "MathComputer",
            "details": operation_details,
            "computation_timestamp": datetime.now().isoformat(),
            "input_numbers": numbers,
            "input_count": len(numbers),
            "additional_params": kwargs
        }
        
        return json.dumps(computation_result, indent=2, ensure_ascii=False)
        
    except Exception as e:
        return f"Error during advanced computation: {str(e)}"

if __name__ == "__main__":
    # Start the MCP server with stdio transport
    mcp.run(transport="stdio")