"""
Secure MCP tool wrapper for CrewAI.

Wraps CrewAI's MCP tool calls with cryptographic message signing,
nonce-based replay protection, and tool definition pinning.
"""

import json
import time
import hashlib
import secrets
from typing import Any, Dict, List, Optional, Set

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import (
    decode_dss_signature,
    encode_dss_signature,
)
import base64


class KeyManager:
    """Manages cryptographic key pairs for tool call signing."""

    def __init__(
        self,
        private_key: Optional[ec.EllipticCurvePrivateKey] = None,
    ):
        if private_key is not None:
            self._private_key = private_key
        else:
            self._private_key = ec.generate_private_key(ec.SECP256R1())
        self._public_key = self._private_key.public_key()

    @property
    def private_key(self) -> ec.EllipticCurvePrivateKey:
        return self._private_key

    @property
    def public_key(self) -> ec.EllipticCurvePublicKey:
        return self._public_key

    def public_key_pem(self) -> str:
        """Export the public key in PEM format."""
        return self._public_key.public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode("utf-8")

    def sign(self, data: bytes) -> str:
        """Sign data and return base64-encoded signature."""
        signature = self._private_key.sign(data, ec.ECDSA(hashes.SHA256()))
        return base64.b64encode(signature).decode("utf-8")

    def verify(self, data: bytes, signature_b64: str) -> bool:
        """Verify a base64-encoded signature against data."""
        try:
            signature = base64.b64decode(signature_b64)
            self._public_key.verify(signature, data, ec.ECDSA(hashes.SHA256()))
            return True
        except Exception:
            return False


class NonceTracker:
    """Tracks nonces to prevent replay attacks."""

    def __init__(self, max_age_seconds: int = 300):
        self._used_nonces: Dict[str, float] = {}
        self._max_age = max_age_seconds

    def generate(self) -> str:
        """Generate a unique nonce."""
        nonce = secrets.token_hex(16)
        self._used_nonces[nonce] = time.time()
        self._cleanup()
        return nonce

    def check_and_consume(self, nonce: str) -> bool:
        """
        Check if a nonce is valid (not previously used).
        Returns True if valid, False if replayed.
        Consumes the nonce on success.
        """
        self._cleanup()
        if nonce in self._used_nonces:
            return False
        self._used_nonces[nonce] = time.time()
        return True

    def is_known(self, nonce: str) -> bool:
        """Check if a nonce was generated by this tracker."""
        return nonce in self._used_nonces

    def _cleanup(self):
        """Remove expired nonces."""
        cutoff = time.time() - self._max_age
        expired = [n for n, t in self._used_nonces.items() if t < cutoff]
        for n in expired:
            del self._used_nonces[n]


class ToolPinner:
    """
    Pins tool definitions at discovery time and detects mutations.

    When an MCP server advertises a tool with a given name and schema,
    we hash that definition. If the definition changes on a subsequent
    call, we raise an error -- this defends against 'rug pull' attacks
    where a server changes tool behaviour after initial trust.
    """

    def __init__(self):
        self._pins: Dict[str, str] = {}

    def pin(self, tool_name: str, tool_definition: Dict[str, Any]) -> str:
        """
        Pin a tool definition. Returns the definition hash.
        """
        canonical = self._canonicalize(tool_definition)
        definition_hash = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
        self._pins[tool_name] = definition_hash
        return definition_hash

    def verify(self, tool_name: str, tool_definition: Dict[str, Any]) -> bool:
        """
        Verify a tool definition matches its pin.
        Returns True if matched or not previously pinned.
        """
        if tool_name not in self._pins:
            return True
        canonical = self._canonicalize(tool_definition)
        current_hash = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
        return current_hash == self._pins[tool_name]

    def get_pin(self, tool_name: str) -> Optional[str]:
        """Get the pinned hash for a tool."""
        return self._pins.get(tool_name)

    def is_pinned(self, tool_name: str) -> bool:
        """Check if a tool has been pinned."""
        return tool_name in self._pins

    @staticmethod
    def _canonicalize(obj: Any) -> str:
        """Produce a canonical JSON string for hashing."""
        return json.dumps(obj, sort_keys=True, separators=(",", ":"))


class SecureMCPTool:
    """
    Wraps CrewAI MCP tool calls with cryptographic security.

    Provides:
    - Cryptographic signing of every tool call
    - Nonce-based replay protection
    - Tool definition pinning (rug-pull detection)
    - Structured audit trail

    Usage with CrewAI's MCPServerAdapter:

        from crewai_mcps import SecureMCPTool

        secure = SecureMCPTool()

        with MCPServerAdapter(server_params) as tools:
            secured_tools = secure.wrap_tools(tools)
            agent = Agent(tools=secured_tools, ...)
    """

    def __init__(
        self,
        key_manager: Optional[KeyManager] = None,
        nonce_ttl: int = 300,
        audit_trail: Optional[Any] = None,
    ):
        self.keys = key_manager or KeyManager()
        self.nonces = NonceTracker(max_age_seconds=nonce_ttl)
        self.pinner = ToolPinner()
        self._audit = audit_trail

    def wrap_tools(self, tools: List[Any]) -> List[Any]:
        """
        Wrap a list of CrewAI MCP tools with security.

        Each tool's definition is pinned at discovery time,
        and every subsequent call is signed with a nonce.

        Args:
            tools: List of tools from MCPServerAdapter

        Returns:
            List of wrapped tools with signing enabled
        """
        from crewai_mcps.middleware import SecureToolProxy

        wrapped = []
        for tool in tools:
            definition = self._extract_definition(tool)
            self.pinner.pin(tool.name, definition)
            proxy = SecureToolProxy(
                original_tool=tool,
                secure_mcp=self,
            )
            wrapped.append(proxy)
        return wrapped

    def sign_call(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, str]:
        """
        Sign a tool call and return the security envelope.

        Args:
            tool_name: Name of the tool being called
            arguments: The arguments being passed to the tool

        Returns:
            Dict with nonce, timestamp, and signature
        """
        nonce = self.nonces.generate()
        timestamp = str(time.time())

        payload = self._build_payload(tool_name, arguments, nonce, timestamp)
        canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        signature = self.keys.sign(canonical.encode("utf-8"))

        envelope = {
            "nonce": nonce,
            "timestamp": timestamp,
            "signature": signature,
            "tool_name": tool_name,
        }

        if self._audit:
            self._audit.log_call(tool_name, arguments, envelope)

        return envelope

    def verify_response(
        self,
        tool_name: str,
        response: Any,
        request_envelope: Dict[str, str],
    ) -> Dict[str, Any]:
        """
        Verify and record a tool response.

        Args:
            tool_name: Name of the tool that was called
            response: The response from the tool
            request_envelope: The security envelope from sign_call

        Returns:
            Verification result dict
        """
        result = {
            "tool_name": tool_name,
            "nonce": request_envelope.get("nonce"),
            "timestamp": request_envelope.get("timestamp"),
            "response_received": True,
            "verified": True,
        }

        if self._audit:
            self._audit.log_response(tool_name, response, result)

        return result

    def check_tool_integrity(self, tool: Any) -> bool:
        """
        Check if a tool's definition still matches its pin.

        Args:
            tool: A CrewAI tool object

        Returns:
            True if the tool definition is unchanged
        """
        definition = self._extract_definition(tool)
        return self.pinner.verify(tool.name, definition)

    def _extract_definition(self, tool: Any) -> Dict[str, Any]:
        """Extract a hashable definition from a CrewAI tool."""
        definition: Dict[str, Any] = {
            "name": getattr(tool, "name", str(tool)),
            "description": getattr(tool, "description", ""),
        }
        if hasattr(tool, "args_schema") and tool.args_schema is not None:
            try:
                definition["schema"] = tool.args_schema.model_json_schema()
            except (AttributeError, TypeError):
                definition["schema"] = str(tool.args_schema)
        return definition

    def _build_payload(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        nonce: str,
        timestamp: str,
    ) -> Dict[str, Any]:
        """Build the canonical payload for signing."""
        return {
            "tool_name": tool_name,
            "arguments": arguments,
            "nonce": nonce,
            "timestamp": timestamp,
        }
