"""
MCP Tool Call Receipts: Audit trail for all MCP tool invocations.

These receipts provide full traceability for MCP tool usage, addressing
IDEsaster-class vulnerabilities where MCP servers become attack vectors.

Receipt types:
- MCPToolCallRequestReceipt: Initial request to invoke a tool
- MCPToolCallVerdictReceipt: Gate decision (ALLOW/DENY/RATE_LIMITED)
- MCPToolCallOutcomeReceipt: Final result of tool execution

Schema: mcp_tool.v1
"""
from __future__ import annotations

import hashlib
import json
import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Literal, Optional

from receipts.compat.pyd import BaseModel, Field
from receipts.base import BaseReceipt, Domain

import logging

logger = logging.getLogger(__name__)


# ============================================================================
# Enums
# ============================================================================

class MCPToolDecision(str, Enum):
    """Decision from MCP Tool Gate."""
    ALLOW = "ALLOW"
    DENY = "DENY"
    RATE_LIMITED = "RATE_LIMITED"


class MCPToolDenyReason(str, Enum):
    """Reason codes for MCP tool denials."""
    TOOL_NOT_FOUND = "tool_not_found"
    TOOL_NOT_ATTESTED = "tool_not_attested"
    ATTESTATION_INVALID = "attestation_invalid"
    ATTESTATION_EXPIRED = "attestation_expired"
    NOT_IN_ALLOWLIST = "not_in_allowlist"
    RATE_LIMITED = "rate_limited"
    SCOPE_MISMATCH = "scope_mismatch"
    PRINCIPAL_NOT_AUTHORIZED = "principal_not_authorized"
    KILL_SWITCH_ACTIVE = "kill_switch_active"


class MCPToolOutcome(str, Enum):
    """Outcome of MCP tool execution."""
    SUCCESS = "success"
    ERROR = "error"
    TIMEOUT = "timeout"
    CANCELLED = "cancelled"


# ============================================================================
# Helper functions
# ============================================================================

def compute_args_hash(args: Optional[Dict[str, Any]]) -> str:
    """Compute deterministic hash of tool arguments."""
    if not args:
        return "sha256:empty"

    canonical = json.dumps(args, sort_keys=True, separators=(",", ":"))
    h = hashlib.sha256(canonical.encode()).hexdigest()[:16]
    return f"sha256:{h}"


def compute_result_hash(result: Any) -> str:
    """Compute hash of tool result."""
    if result is None:
        return "sha256:null"

    try:
        canonical = json.dumps(result, sort_keys=True, separators=(",", ":"), default=str)
        h = hashlib.sha256(canonical.encode()).hexdigest()[:16]
        return f"sha256:{h}"
    except Exception:
        return "sha256:unhashable"


# ============================================================================
# Receipt: MCP Tool Call Request
# ============================================================================

class MCPToolCallRequestReceipt(BaseReceipt):
    """
    Receipt emitted when an MCP tool call is requested.

    This captures the initial request before gate evaluation.
    Links to subsequent verdict and outcome receipts.

    Schema: mcp_tool.v1
    """

    receipt_type: str = Field(default="MCPToolCallRequestReceipt")
    domain: str = Field(default=Domain.GOVERNANCE.value)

    # Request identity
    call_id: str = Field(
        description="Unique ID for this tool call"
    )
    tool_name: str = Field(
        description="Name of the tool being called (e.g., 'loom.execute_episode')"
    )

    # Caller identity
    principal_id: str = Field(
        description="Who is calling the tool (agent/user ID)"
    )
    principal_type: str = Field(
        default="agent",
        description="Type of principal: 'agent', 'user', 'service'"
    )

    # Request details
    args_hash: str = Field(
        description="Hash of tool arguments (privacy-preserving)"
    )
    args_keys: List[str] = Field(
        default_factory=list,
        description="List of argument keys (not values)"
    )

    # Context
    episode_id: Optional[str] = Field(
        default=None,
        description="Episode this call belongs to"
    )
    session_id: Optional[str] = Field(
        default=None,
        description="MCP session ID"
    )

    # Attestation reference
    attestation_id: Optional[str] = Field(
        default=None,
        description="ID of tool attestation being used (if any)"
    )

    @classmethod
    def create(
        cls,
        tool_name: str,
        principal_id: str,
        args: Optional[Dict[str, Any]] = None,
        principal_type: str = "agent",
        episode_id: Optional[str] = None,
        session_id: Optional[str] = None,
        attestation_id: Optional[str] = None,
    ) -> "MCPToolCallRequestReceipt":
        """Create a new request receipt."""
        now = datetime.now(timezone.utc)
        call_id = f"mcp_call_{now.strftime('%Y%m%dT%H%M%S')}_{uuid.uuid4().hex[:8]}"

        return cls(
            receipt_id=f"rcpt_mcp_req_{call_id}",
            ts=now,
            call_id=call_id,
            tool_name=tool_name,
            principal_id=principal_id,
            principal_type=principal_type,
            args_hash=compute_args_hash(args),
            args_keys=list(args.keys()) if args else [],
            episode_id=episode_id,
            session_id=session_id,
            attestation_id=attestation_id,
        )


# ============================================================================
# Receipt: MCP Tool Call Verdict
# ============================================================================

class MCPToolCallVerdictReceipt(BaseReceipt):
    """
    Receipt emitted by MCP Tool Gate with its decision.

    Links to the request receipt and captures:
    - ALLOW/DENY/RATE_LIMITED decision
    - Reason for denial (if applicable)
    - Gate checks performed

    Schema: mcp_tool.v1
    """

    receipt_type: str = Field(default="MCPToolCallVerdictReceipt")
    domain: str = Field(default=Domain.GOVERNANCE.value)

    # Link to request
    request_receipt_id: str = Field(
        description="FK to MCPToolCallRequestReceipt"
    )
    call_id: str = Field(
        description="Tool call ID (from request)"
    )

    # Decision
    verdict: MCPToolDecision = Field(
        description="Gate decision: ALLOW, DENY, or RATE_LIMITED"
    )
    deny_reason: Optional[MCPToolDenyReason] = Field(
        default=None,
        description="Reason code if denied"
    )
    deny_detail: Optional[str] = Field(
        default=None,
        description="Human-readable denial explanation"
    )

    # Gate checks
    attestation_verified: bool = Field(
        default=False,
        description="Whether tool attestation was verified"
    )
    allowlist_checked: bool = Field(
        default=False,
        description="Whether allowlist was checked"
    )
    rate_limit_checked: bool = Field(
        default=False,
        description="Whether rate limit was checked"
    )

    # Rate limit info (if applicable)
    rate_limit_remaining: Optional[int] = Field(
        default=None,
        description="Remaining calls in current window"
    )
    rate_limit_reset_at: Optional[datetime] = Field(
        default=None,
        description="When rate limit resets"
    )

    # Timing
    latency_ms: float = Field(
        default=0.0,
        description="Gate evaluation time in milliseconds"
    )

    @classmethod
    def create_allow(
        cls,
        request_receipt_id: str,
        call_id: str,
        attestation_verified: bool = False,
        allowlist_checked: bool = False,
        rate_limit_checked: bool = False,
        rate_limit_remaining: Optional[int] = None,
        latency_ms: float = 0.0,
    ) -> "MCPToolCallVerdictReceipt":
        """Create an ALLOW verdict."""
        return cls(
            receipt_id=f"rcpt_mcp_verdict_{call_id}",
            ts=datetime.now(timezone.utc),
            request_receipt_id=request_receipt_id,
            call_id=call_id,
            verdict=MCPToolDecision.ALLOW,
            attestation_verified=attestation_verified,
            allowlist_checked=allowlist_checked,
            rate_limit_checked=rate_limit_checked,
            rate_limit_remaining=rate_limit_remaining,
            latency_ms=latency_ms,
        )

    @classmethod
    def create_deny(
        cls,
        request_receipt_id: str,
        call_id: str,
        reason: MCPToolDenyReason,
        detail: Optional[str] = None,
        attestation_verified: bool = False,
        allowlist_checked: bool = False,
        rate_limit_checked: bool = False,
        latency_ms: float = 0.0,
    ) -> "MCPToolCallVerdictReceipt":
        """Create a DENY verdict."""
        return cls(
            receipt_id=f"rcpt_mcp_verdict_{call_id}",
            ts=datetime.now(timezone.utc),
            request_receipt_id=request_receipt_id,
            call_id=call_id,
            verdict=MCPToolDecision.DENY,
            deny_reason=reason,
            deny_detail=detail,
            attestation_verified=attestation_verified,
            allowlist_checked=allowlist_checked,
            rate_limit_checked=rate_limit_checked,
            latency_ms=latency_ms,
        )

    @classmethod
    def create_rate_limited(
        cls,
        request_receipt_id: str,
        call_id: str,
        rate_limit_reset_at: Optional[datetime] = None,
        latency_ms: float = 0.0,
    ) -> "MCPToolCallVerdictReceipt":
        """Create a RATE_LIMITED verdict."""
        return cls(
            receipt_id=f"rcpt_mcp_verdict_{call_id}",
            ts=datetime.now(timezone.utc),
            request_receipt_id=request_receipt_id,
            call_id=call_id,
            verdict=MCPToolDecision.RATE_LIMITED,
            deny_reason=MCPToolDenyReason.RATE_LIMITED,
            rate_limit_checked=True,
            rate_limit_remaining=0,
            rate_limit_reset_at=rate_limit_reset_at,
            latency_ms=latency_ms,
        )


# ============================================================================
# Receipt: MCP Tool Call Outcome
# ============================================================================

class MCPToolCallOutcomeReceipt(BaseReceipt):
    """
    Receipt emitted after tool execution completes.

    Captures the final outcome (success/error/timeout) and
    links back to request and verdict receipts.

    Schema: mcp_tool.v1
    """

    receipt_type: str = Field(default="MCPToolCallOutcomeReceipt")
    domain: str = Field(default=Domain.GOVERNANCE.value)

    # Links
    request_receipt_id: str = Field(
        description="FK to MCPToolCallRequestReceipt"
    )
    verdict_receipt_id: str = Field(
        description="FK to MCPToolCallVerdictReceipt"
    )
    call_id: str = Field(
        description="Tool call ID"
    )

    # Outcome
    outcome: MCPToolOutcome = Field(
        description="Execution outcome"
    )
    result_hash: Optional[str] = Field(
        default=None,
        description="Hash of result (for verification)"
    )
    error_code: Optional[str] = Field(
        default=None,
        description="Error code if outcome is ERROR"
    )
    error_message: Optional[str] = Field(
        default=None,
        description="Error message if outcome is ERROR"
    )

    # Timing
    execution_ms: float = Field(
        default=0.0,
        description="Tool execution time in milliseconds"
    )
    total_ms: float = Field(
        default=0.0,
        description="Total time including gate overhead"
    )

    @classmethod
    def create_success(
        cls,
        request_receipt_id: str,
        verdict_receipt_id: str,
        call_id: str,
        result: Any = None,
        execution_ms: float = 0.0,
        total_ms: float = 0.0,
    ) -> "MCPToolCallOutcomeReceipt":
        """Create a SUCCESS outcome."""
        return cls(
            receipt_id=f"rcpt_mcp_outcome_{call_id}",
            ts=datetime.now(timezone.utc),
            request_receipt_id=request_receipt_id,
            verdict_receipt_id=verdict_receipt_id,
            call_id=call_id,
            outcome=MCPToolOutcome.SUCCESS,
            result_hash=compute_result_hash(result),
            execution_ms=execution_ms,
            total_ms=total_ms,
        )

    @classmethod
    def create_error(
        cls,
        request_receipt_id: str,
        verdict_receipt_id: str,
        call_id: str,
        error_code: str,
        error_message: str,
        execution_ms: float = 0.0,
        total_ms: float = 0.0,
    ) -> "MCPToolCallOutcomeReceipt":
        """Create an ERROR outcome."""
        return cls(
            receipt_id=f"rcpt_mcp_outcome_{call_id}",
            ts=datetime.now(timezone.utc),
            request_receipt_id=request_receipt_id,
            verdict_receipt_id=verdict_receipt_id,
            call_id=call_id,
            outcome=MCPToolOutcome.ERROR,
            error_code=error_code,
            error_message=error_message,
            execution_ms=execution_ms,
            total_ms=total_ms,
        )


# ============================================================================
# Tool Attestation Model
# ============================================================================

class ToolAttestation(BaseModel):
    """
    Tool attestation record.

    Attestations verify tool identity and authorize its use.
    Similar to a code signing certificate for executables.

    Fields:
    - tool_id: Unique tool identifier
    - version: Semantic version
    - publisher: Who published the tool
    - signature: Cryptographic signature
    - allowed_scopes: What the tool is allowed to do
    - expires_at: When attestation expires
    """

    attestation_id: str = Field(
        description="Unique attestation ID"
    )
    tool_id: str = Field(
        description="Tool identifier (e.g., 'loom.execute_episode')"
    )
    version: str = Field(
        default="1.0.0",
        description="Tool version (semver)"
    )
    publisher: str = Field(
        description="Publisher identity (e.g., 'loom-labs')"
    )

    # Scopes and capabilities
    allowed_scopes: List[str] = Field(
        default_factory=list,
        description="Scopes this tool can access (e.g., ['fs.read', 'http.fetch'])"
    )
    max_execution_ms: Optional[int] = Field(
        default=None,
        description="Maximum execution time (soft limit)"
    )

    # Signature
    signature: Optional[str] = Field(
        default=None,
        description="Ed25519 signature of attestation fields"
    )
    signature_algorithm: str = Field(
        default="ed25519",
        description="Signature algorithm"
    )

    # Validity
    issued_at: datetime = Field(
        default_factory=lambda: datetime.now(timezone.utc)
    )
    expires_at: Optional[datetime] = Field(
        default=None,
        description="When attestation expires"
    )

    def is_expired(self) -> bool:
        """Check if attestation has expired."""
        if self.expires_at is None:
            return False
        return datetime.now(timezone.utc) >= self.expires_at

    def covers_scope(self, scope: str) -> bool:
        """Check if attestation covers a scope."""
        if not self.allowed_scopes:
            return True  # No restrictions

        for allowed in self.allowed_scopes:
            if allowed == "*":
                return True
            if allowed == scope:
                return True
            # Check wildcard prefix (e.g., "fs.*" covers "fs.read")
            if allowed.endswith(".*"):
                prefix = allowed[:-2]
                if scope.startswith(prefix + "."):
                    return True

        return False

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict for serialization."""
        return {
            "attestation_id": self.attestation_id,
            "tool_id": self.tool_id,
            "version": self.version,
            "publisher": self.publisher,
            "allowed_scopes": self.allowed_scopes,
            "max_execution_ms": self.max_execution_ms,
            "signature": self.signature,
            "signature_algorithm": self.signature_algorithm,
            "issued_at": self.issued_at.isoformat() if self.issued_at else None,
            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ToolAttestation":
        """Create from dict."""
        # Handle datetime fields
        if "issued_at" in data and isinstance(data["issued_at"], str):
            data["issued_at"] = datetime.fromisoformat(data["issued_at"].replace("Z", "+00:00"))
        if "expires_at" in data and isinstance(data["expires_at"], str):
            data["expires_at"] = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00"))

        return cls(**data)


__all__ = [
    # Enums
    "MCPToolDecision",
    "MCPToolDenyReason",
    "MCPToolOutcome",
    # Receipts
    "MCPToolCallRequestReceipt",
    "MCPToolCallVerdictReceipt",
    "MCPToolCallOutcomeReceipt",
    # Models
    "ToolAttestation",
    # Helpers
    "compute_args_hash",
    "compute_result_hash",
]
