"""
Handler signature validation for tool registration.

Validates that handler function signatures match the tool's input schema,
catching mismatches at registration time rather than runtime.

Supports two handler patterns:
1. Standard args dict pattern: `async def handler(args: dict) -> dict`
2. Typed parameters pattern: `async def handler(to: str, subject: str) -> dict`

For the typed pattern, validates that:
- All required schema fields have corresponding required parameters
- Parameter types match schema types
- No unexpected parameters that aren't in the schema
"""
import inspect
import logging
from dataclasses import dataclass, field
from typing import Any, Callable, Optional, Union, get_type_hints, get_origin, get_args

logger = logging.getLogger(__name__)


@dataclass
class SignatureValidationError:
    """Represents a single validation error."""
    field: str
    message: str
    severity: str = "error"  # error or warning


@dataclass
class SignatureValidationResult:
    """Result of validating a handler signature against a schema."""
    is_valid: bool
    errors: list[SignatureValidationError] = field(default_factory=list)
    warnings: list[SignatureValidationError] = field(default_factory=list)
    handler_pattern: str = "unknown"  # "args_dict" or "typed_params"

    def raise_if_invalid(self, tool_name: str) -> None:
        """Raise TypeError if validation failed."""
        if not self.is_valid:
            error_messages = [f"  - {e.field}: {e.message}" for e in self.errors]
            raise TypeError(
                f"Handler signature mismatch for tool '{tool_name}':\n"
                + "\n".join(error_messages)
            )


class HandlerSignatureValidator:
    """
    Validates handler function signatures against tool schemas.

    Detects handler pattern and validates accordingly:
    - args dict pattern: Just validates single dict parameter
    - typed params pattern: Validates parameters match schema fields
    """

    # Mapping from Python types to JSON Schema types
    PYTHON_TO_JSON_SCHEMA = {
        str: "string",
        int: "integer",
        float: "number",
        bool: "boolean",
        list: "array",
        dict: "object",
        type(None): "null",
    }

    # Mapping from JSON Schema types to compatible Python types
    JSON_SCHEMA_TO_PYTHON = {
        "string": (str,),
        "integer": (int,),
        "number": (int, float),  # JSON number can be int or float
        "boolean": (bool,),
        "array": (list,),
        "object": (dict,),
        "null": (type(None),),
    }

    def __init__(self, strict: bool = True):
        """
        Initialize validator.

        Args:
            strict: If True, warnings are treated as errors
        """
        self.strict = strict

    def validate(
        self,
        handler: Callable,
        schema: dict[str, Any],
        tool_name: str,
    ) -> SignatureValidationResult:
        """
        Validate a handler's signature against a schema.

        Args:
            handler: The async handler function
            schema: Input schema (Python types or JSON Schema)
            tool_name: Name of the tool (for error messages)

        Returns:
            SignatureValidationResult with validation details
        """
        errors: list[SignatureValidationError] = []
        warnings: list[SignatureValidationError] = []

        # Get handler signature
        try:
            sig = inspect.signature(handler)
        except (ValueError, TypeError) as e:
            return SignatureValidationResult(
                is_valid=False,
                errors=[SignatureValidationError(
                    field="handler",
                    message=f"Could not inspect handler signature: {e}",
                )],
            )

        # Get type hints if available
        try:
            type_hints = get_type_hints(handler)
        except Exception:
            type_hints = {}

        # Determine handler pattern
        params = list(sig.parameters.values())
        pattern = self._detect_pattern(params, type_hints)

        if pattern == "args_dict":
            # Standard args dict pattern - minimal validation
            validation_errors = self._validate_args_dict_pattern(params, type_hints)
            errors.extend(validation_errors)
        elif pattern == "typed_params":
            # Typed parameters pattern - full validation
            param_errors, param_warnings = self._validate_typed_params_pattern(
                params, type_hints, schema, tool_name
            )
            errors.extend(param_errors)
            warnings.extend(param_warnings)
        else:
            # Unknown pattern
            errors.append(SignatureValidationError(
                field="handler",
                message="Handler must accept either a single 'args: dict' parameter "
                        "or typed parameters matching the schema",
            ))

        # Validate return type if annotated
        return_errors = self._validate_return_type(type_hints, tool_name)
        errors.extend(return_errors)

        # Determine validity
        is_valid = len(errors) == 0
        if self.strict and warnings:
            is_valid = False
            errors.extend(warnings)
            warnings = []

        return SignatureValidationResult(
            is_valid=is_valid,
            errors=errors,
            warnings=warnings,
            handler_pattern=pattern,
        )

    def _detect_pattern(
        self,
        params: list[inspect.Parameter],
        type_hints: dict[str, Any],
    ) -> str:
        """
        Detect which handler pattern is being used.

        Returns:
            "args_dict" - single dict parameter named 'args'
            "typed_params" - individual typed parameters
            "unknown" - unrecognized pattern
        """
        # Filter out *args and **kwargs
        positional_params = [
            p for p in params
            if p.kind not in (
                inspect.Parameter.VAR_POSITIONAL,
                inspect.Parameter.VAR_KEYWORD,
            )
        ]

        if len(positional_params) == 0:
            return "unknown"

        if len(positional_params) == 1:
            param = positional_params[0]
            param_type = type_hints.get(param.name)

            # Check if it's the args dict pattern
            if param.name == "args":
                if param_type is None or param_type is dict or self._is_dict_type(param_type):
                    return "args_dict"

            # Single param with dict type but different name - still args dict pattern
            if param_type is dict or self._is_dict_type(param_type):
                return "args_dict"

        # Multiple params or single non-dict param = typed params pattern
        return "typed_params"

    def _is_dict_type(self, type_hint: Any) -> bool:
        """Check if a type hint is a dict type."""
        if type_hint is dict:
            return True
        origin = get_origin(type_hint)
        return origin is dict

    def _validate_args_dict_pattern(
        self,
        params: list[inspect.Parameter],
        type_hints: dict[str, Any],
    ) -> list[SignatureValidationError]:
        """Validate the args dict handler pattern."""
        errors = []

        positional_params = [
            p for p in params
            if p.kind not in (
                inspect.Parameter.VAR_POSITIONAL,
                inspect.Parameter.VAR_KEYWORD,
            )
        ]

        if len(positional_params) != 1:
            errors.append(SignatureValidationError(
                field="handler",
                message=f"Args dict pattern requires exactly 1 parameter, got {len(positional_params)}",
            ))
            return errors

        param = positional_params[0]
        param_type = type_hints.get(param.name)

        # Validate type annotation if present
        if param_type is not None:
            if not (param_type is dict or self._is_dict_type(param_type)):
                errors.append(SignatureValidationError(
                    field=param.name,
                    message=f"Parameter should be 'dict', got '{param_type}'",
                ))

        return errors

    def _validate_typed_params_pattern(
        self,
        params: list[inspect.Parameter],
        type_hints: dict[str, Any],
        schema: dict[str, Any],
        tool_name: str,
    ) -> tuple[list[SignatureValidationError], list[SignatureValidationError]]:
        """Validate the typed parameters handler pattern."""
        errors = []
        warnings = []

        # Normalize schema to get fields and their types
        schema_fields = self._normalize_schema(schema)

        # Get positional params (excluding *args, **kwargs)
        positional_params = [
            p for p in params
            if p.kind not in (
                inspect.Parameter.VAR_POSITIONAL,
                inspect.Parameter.VAR_KEYWORD,
            )
        ]

        param_names = {p.name for p in positional_params}
        schema_field_names = set(schema_fields.keys())

        # Check for required schema fields missing from handler
        required_fields = {
            name for name, info in schema_fields.items()
            if info.get("required", True)
        }
        missing_required = required_fields - param_names

        # Check if handler has **kwargs to catch extra fields
        has_kwargs = any(
            p.kind == inspect.Parameter.VAR_KEYWORD for p in params
        )

        if missing_required and not has_kwargs:
            for field_name in sorted(missing_required):
                errors.append(SignatureValidationError(
                    field=field_name,
                    message=f"Required schema field '{field_name}' not found in handler parameters",
                ))

        # Check for handler params not in schema
        extra_params = param_names - schema_field_names
        for param_name in sorted(extra_params):
            warnings.append(SignatureValidationError(
                field=param_name,
                message=f"Handler parameter '{param_name}' not found in schema",
                severity="warning",
            ))

        # Validate types for matching parameters
        for param in positional_params:
            if param.name not in schema_fields:
                continue

            schema_info = schema_fields[param.name]
            param_type = type_hints.get(param.name)

            if param_type is None:
                # No type annotation - can't validate, just warn
                warnings.append(SignatureValidationError(
                    field=param.name,
                    message=f"Parameter '{param.name}' has no type annotation",
                    severity="warning",
                ))
                continue

            # Validate type compatibility
            type_error = self._check_type_compatibility(
                param.name, param_type, schema_info
            )
            if type_error:
                errors.append(type_error)

            # Check optional vs required
            is_param_optional = (
                param.default is not inspect.Parameter.empty
                or self._is_optional_type(param_type)
            )
            is_schema_required = schema_info.get("required", True)

            if is_schema_required and is_param_optional:
                warnings.append(SignatureValidationError(
                    field=param.name,
                    message=f"Schema requires '{param.name}' but handler parameter has default value",
                    severity="warning",
                ))

        return errors, warnings

    def _normalize_schema(self, schema: dict[str, Any]) -> dict[str, dict[str, Any]]:
        """
        Normalize schema to consistent format.

        Returns dict mapping field names to their info:
        {
            "field_name": {
                "type": "string",  # JSON Schema type
                "required": True,
                "python_type": str,  # Original Python type if available
            }
        }
        """
        result = {}

        # Check if it's JSON Schema format
        if schema.get("type") == "object" and "properties" in schema:
            properties = schema.get("properties", {})
            required_fields = set(schema.get("required", []))

            for name, prop_schema in properties.items():
                result[name] = {
                    "type": prop_schema.get("type", "string"),
                    "required": name in required_fields,
                    "schema": prop_schema,
                }
        else:
            # Simple Python type schema: {"name": str, "count": int}
            for name, type_hint in schema.items():
                is_optional = self._is_optional_type(type_hint)
                json_type = self._python_type_to_json_type(type_hint)

                result[name] = {
                    "type": json_type,
                    "required": not is_optional,
                    "python_type": type_hint,  # Keep original for Union comparisons
                }

        return result

    def _python_type_to_json_type(self, type_hint: Any) -> str:
        """Convert Python type to JSON Schema type string."""
        # Handle None
        if type_hint is None or type_hint is type(None):
            return "null"

        # Handle basic types
        if type_hint in self.PYTHON_TO_JSON_SCHEMA:
            return self.PYTHON_TO_JSON_SCHEMA[type_hint]

        # Handle typing module types
        origin = get_origin(type_hint)
        args = get_args(type_hint)

        # Handle Optional/Union
        if origin is Union:
            non_none_args = [a for a in args if a is not type(None)]
            if non_none_args:
                return self._python_type_to_json_type(non_none_args[0])
            return "null"

        # Handle List
        if origin is list:
            return "array"

        # Handle Dict
        if origin is dict:
            return "object"

        # Fallback
        return "string"

    def _is_optional_type(self, type_hint: Any) -> bool:
        """Check if a type hint is Optional (Union with None)."""
        origin = get_origin(type_hint)
        args = get_args(type_hint)

        if origin is Union:
            return type(None) in args
        return False

    def _check_type_compatibility(
        self,
        param_name: str,
        param_type: Any,
        schema_info: dict[str, Any],
    ) -> Optional[SignatureValidationError]:
        """
        Check if a parameter type is compatible with schema type.

        Returns error if incompatible, None if compatible.
        """
        schema_type = schema_info.get("type", "string")
        schema_python_type = schema_info.get("python_type")

        # Get the base type (unwrap Optional)
        base_type = self._unwrap_optional(param_type)

        # If schema has original Python type (e.g., Union), compare directly
        if schema_python_type is not None:
            # Handle Union comparisons - check if the types are equivalent
            schema_origin = get_origin(schema_python_type)
            param_origin = get_origin(base_type)

            if schema_origin is Union and param_origin is Union:
                # Both are Union types - compare their args
                schema_args = set(get_args(self._unwrap_optional(schema_python_type)))
                param_args = set(get_args(base_type))
                # Types are compatible if they have the same type arguments
                if schema_args == param_args:
                    return None

            # Direct Python type match
            unwrapped_schema = self._unwrap_optional(schema_python_type)
            if base_type == unwrapped_schema:
                return None

            # Check if the param type is one of the types in a Union schema
            if schema_origin is Union:
                schema_args = get_args(self._unwrap_optional(schema_python_type))
                if base_type in schema_args:
                    return None

        # Get expected Python types for this JSON Schema type
        expected_types = self.JSON_SCHEMA_TO_PYTHON.get(schema_type, (str,))

        # Check if it's a generic (List, Dict, etc.)
        origin = get_origin(base_type)

        # Handle Union types in parameter
        if origin is Union:
            # For Union types, check if all non-None args are compatible
            union_args = [a for a in get_args(base_type) if a is not type(None)]
            all_compatible = all(
                arg in expected_types or get_origin(arg) in expected_types
                for arg in union_args
            )
            if all_compatible:
                return None

        if origin is not None:
            # For generics, check the origin type
            if origin not in expected_types and origin is not Union:
                return SignatureValidationError(
                    field=param_name,
                    message=f"Type mismatch: parameter is '{base_type}' but schema expects '{schema_type}'",
                )
        elif base_type not in expected_types and base_type is not Any:
            return SignatureValidationError(
                field=param_name,
                message=f"Type mismatch: parameter is '{base_type.__name__ if hasattr(base_type, '__name__') else base_type}' "
                        f"but schema expects '{schema_type}'",
            )

        return None

    def _unwrap_optional(self, type_hint: Any) -> Any:
        """Unwrap Optional[X] to get X."""
        origin = get_origin(type_hint)
        args = get_args(type_hint)

        if origin is Union:
            non_none_args = [a for a in args if a is not type(None)]
            if len(non_none_args) == 1:
                return non_none_args[0]
        return type_hint

    def _validate_return_type(
        self,
        type_hints: dict[str, Any],
        tool_name: str,
    ) -> list[SignatureValidationError]:
        """Validate the handler's return type annotation."""
        errors = []

        return_type = type_hints.get("return")
        if return_type is None:
            # No return annotation - warn but don't error
            return []

        # Check if return type is dict-compatible
        if return_type is dict or self._is_dict_type(return_type):
            return []

        # Check for Any
        if return_type is Any:
            return []

        # Check for None (invalid)
        if return_type is type(None):
            errors.append(SignatureValidationError(
                field="return",
                message="Handler must return a dict, not None",
            ))
            return errors

        errors.append(SignatureValidationError(
            field="return",
            message=f"Handler should return 'dict', got '{return_type}'",
        ))

        return errors


def validate_handler_signature(
    handler: Callable,
    schema: dict[str, Any],
    tool_name: str,
    strict: bool = False,
) -> SignatureValidationResult:
    """
    Convenience function to validate a handler signature.

    Args:
        handler: The async handler function
        schema: Input schema (Python types or JSON Schema)
        tool_name: Name of the tool (for error messages)
        strict: If True, warnings are treated as errors

    Returns:
        SignatureValidationResult with validation details
    """
    validator = HandlerSignatureValidator(strict=strict)
    return validator.validate(handler, schema, tool_name)
