"""
Middleware for intercepting CrewAI MCP tool calls.

Provides pre/post call hooks that add cryptographic signing
and verification to every tool invocation.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from crewai_mcps.secure_tool import SecureMCPTool

logger = logging.getLogger("crewai_mcps")


class SecureToolProxy:
    """
    Proxy that wraps a CrewAI tool with MCPS signing.

    Intercepts _run() calls to:
    1. Verify tool definition integrity (rug-pull detection)
    2. Sign the outgoing call with a nonce
    3. Execute the original tool
    4. Verify and log the response

    The proxy delegates attribute access to the original tool,
    so it works as a drop-in replacement in CrewAI agent configs.
    """

    def __init__(
        self,
        original_tool: Any,
        secure_mcp: "SecureMCPTool",
    ):
        self._original = original_tool
        self._secure = secure_mcp
        self._call_count = 0

    # -- Proxy all attribute access to the original tool --

    @property
    def name(self) -> str:
        return self._original.name

    @property
    def description(self) -> str:
        return self._original.description

    @property
    def args_schema(self):
        return getattr(self._original, "args_schema", None)

    def __getattr__(self, name: str) -> Any:
        """Delegate unknown attribute access to the original tool."""
        return getattr(self._original, name)

    # -- Intercept tool execution --

    def run(self, *args: Any, **kwargs: Any) -> Any:
        """Signed execution via run()."""
        return self._execute_signed(args, kwargs)

    def _run(self, *args: Any, **kwargs: Any) -> Any:
        """Signed execution via _run()."""
        return self._execute_signed(args, kwargs)

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Signed execution via direct call."""
        return self._execute_signed(args, kwargs)

    def _execute_signed(self, args: tuple, kwargs: dict) -> Any:
        """
        Core signed execution flow.

        1. Check tool integrity against pin
        2. Sign the call
        3. Execute original tool
        4. Verify response
        5. Return result
        """
        self._call_count += 1

        # Step 1: Verify tool definition hasn't changed
        if not self._secure.check_tool_integrity(self._original):
            raise ToolIntegrityError(
                f"Tool '{self.name}' definition has changed since discovery. "
                "Possible rug-pull attack detected. Refusing to execute."
            )

        # Step 2: Build arguments dict for signing
        arguments = self._build_arguments(args, kwargs)

        # Step 3: Sign the call
        envelope = self._secure.sign_call(self.name, arguments)

        logger.debug(
            "MCPS: Signed call to '%s' [nonce=%s]",
            self.name,
            envelope["nonce"][:8] + "...",
        )

        # Step 4: Execute the original tool
        try:
            if hasattr(self._original, "run"):
                result = self._original.run(*args, **kwargs)
            elif hasattr(self._original, "_run"):
                result = self._original._run(*args, **kwargs)
            else:
                result = self._original(*args, **kwargs)
        except Exception as e:
            if self._secure._audit:
                self._secure._audit.log_error(self.name, str(e), envelope)
            raise

        # Step 5: Verify and log response
        self._secure.verify_response(self.name, result, envelope)

        logger.debug(
            "MCPS: Verified response from '%s' [nonce=%s]",
            self.name,
            envelope["nonce"][:8] + "...",
        )

        return result

    def _build_arguments(self, args: tuple, kwargs: dict) -> Dict[str, Any]:
        """Build a serializable arguments dict from args/kwargs."""
        arguments: Dict[str, Any] = {}
        if args:
            arguments["_positional"] = [
                self._safe_serialize(a) for a in args
            ]
        for k, v in kwargs.items():
            arguments[k] = self._safe_serialize(v)
        return arguments

    @staticmethod
    def _safe_serialize(value: Any) -> Any:
        """Make a value JSON-serializable."""
        if isinstance(value, (str, int, float, bool, type(None))):
            return value
        if isinstance(value, (list, tuple)):
            return [SecureToolProxy._safe_serialize(v) for v in value]
        if isinstance(value, dict):
            return {
                str(k): SecureToolProxy._safe_serialize(v)
                for k, v in value.items()
            }
        return str(value)


class ToolIntegrityError(Exception):
    """Raised when a tool definition changes after pinning."""
    pass


def install_crewai_hooks(secure_mcp: "SecureMCPTool") -> Callable:
    """
    Install MCPS signing as CrewAI tool hooks.

    This uses CrewAI's built-in hook system to add signing
    to ALL tool calls globally, not just MCP tools.

    Args:
        secure_mcp: A configured SecureMCPTool instance

    Returns:
        Cleanup function to uninstall the hooks
    """
    _envelopes: Dict[str, Dict[str, str]] = {}

    def before_hook(context: Any) -> Optional[bool]:
        """Sign tool calls before execution."""
        tool_name = context.tool_name
        tool_input = dict(context.tool_input) if context.tool_input else {}

        envelope = secure_mcp.sign_call(tool_name, tool_input)
        _envelopes[envelope["nonce"]] = envelope

        logger.debug(
            "MCPS hook: Signed call to '%s' [nonce=%s]",
            tool_name,
            envelope["nonce"][:8] + "...",
        )
        return None  # Allow execution to proceed

    def after_hook(context: Any) -> Optional[str]:
        """Verify tool responses after execution."""
        tool_name = context.tool_name
        # Find the most recent envelope for this tool
        recent = None
        for nonce, env in list(_envelopes.items()):
            if env["tool_name"] == tool_name:
                recent = env
                del _envelopes[nonce]
                break

        if recent:
            secure_mcp.verify_response(
                tool_name, context.tool_result, recent
            )
            logger.debug(
                "MCPS hook: Verified response from '%s'",
                tool_name,
            )

        return None  # Don't modify the result

    # Try to install hooks if crewai is available
    try:
        from crewai.hooks.tool_hooks import (
            register_before_tool_call_hook,
            register_after_tool_call_hook,
            unregister_before_tool_call_hook,
            unregister_after_tool_call_hook,
        )

        register_before_tool_call_hook(before_hook)
        register_after_tool_call_hook(after_hook)

        def cleanup():
            unregister_before_tool_call_hook(before_hook)
            unregister_after_tool_call_hook(after_hook)

        return cleanup

    except ImportError:
        logger.warning(
            "crewai not installed; hook-based signing not available. "
            "Use SecureMCPTool.wrap_tools() instead."
        )
        return lambda: None
