"""
Structured audit trail for MCPS-signed tool calls.

Records every signed call, response, and error with timestamps,
nonces, and tool metadata for compliance and forensic analysis.
"""

import json
import time
import threading
from typing import Any, Callable, Dict, List, Optional
from dataclasses import dataclass, field, asdict
from enum import Enum


class AuditEventType(str, Enum):
    """Types of auditable events."""
    TOOL_CALL = "tool_call"
    TOOL_RESPONSE = "tool_response"
    TOOL_ERROR = "tool_error"
    PIN_CREATED = "pin_created"
    PIN_VERIFIED = "pin_verified"
    PIN_MISMATCH = "pin_mismatch"
    REPLAY_DETECTED = "replay_detected"
    KEY_GENERATED = "key_generated"


@dataclass
class AuditEvent:
    """A single audit trail entry."""
    event_type: str
    tool_name: str
    timestamp: float
    nonce: Optional[str] = None
    details: Dict[str, Any] = field(default_factory=dict)
    sequence: int = 0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to a plain dictionary."""
        return asdict(self)

    def to_json(self) -> str:
        """Serialize to JSON string."""
        return json.dumps(self.to_dict(), sort_keys=True)


class AuditTrail:
    """
    Thread-safe structured audit trail for MCPS operations.

    Records all signing, verification, and security events
    with monotonically increasing sequence numbers.

    Usage:
        audit = AuditTrail()
        secure = SecureMCPTool(audit_trail=audit)

        # After operations...
        for event in audit.events:
            print(event.to_json())

        # Export full trail
        audit.export_json("audit_trail.json")
    """

    def __init__(
        self,
        max_events: int = 10000,
        on_event: Optional[Callable[[AuditEvent], None]] = None,
    ):
        self._events: List[AuditEvent] = []
        self._lock = threading.Lock()
        self._sequence = 0
        self._max_events = max_events
        self._on_event = on_event

    @property
    def events(self) -> List[AuditEvent]:
        """Get a copy of all events."""
        with self._lock:
            return list(self._events)

    @property
    def event_count(self) -> int:
        """Number of recorded events."""
        with self._lock:
            return len(self._events)

    def log_call(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        envelope: Dict[str, str],
    ) -> AuditEvent:
        """Record a signed tool call."""
        return self._record(
            AuditEventType.TOOL_CALL,
            tool_name,
            nonce=envelope.get("nonce"),
            details={
                "argument_keys": list(arguments.keys()),
                "argument_count": len(arguments),
                "signature_prefix": envelope.get("signature", "")[:16] + "...",
            },
        )

    def log_response(
        self,
        tool_name: str,
        response: Any,
        verification: Dict[str, Any],
    ) -> AuditEvent:
        """Record a verified tool response."""
        response_str = str(response)
        return self._record(
            AuditEventType.TOOL_RESPONSE,
            tool_name,
            nonce=verification.get("nonce"),
            details={
                "response_length": len(response_str),
                "response_preview": response_str[:100] + ("..." if len(response_str) > 100 else ""),
                "verified": verification.get("verified", False),
            },
        )

    def log_error(
        self,
        tool_name: str,
        error: str,
        envelope: Optional[Dict[str, str]] = None,
    ) -> AuditEvent:
        """Record a tool call error."""
        return self._record(
            AuditEventType.TOOL_ERROR,
            tool_name,
            nonce=envelope.get("nonce") if envelope else None,
            details={"error": error},
        )

    def log_pin_created(
        self,
        tool_name: str,
        definition_hash: str,
    ) -> AuditEvent:
        """Record tool definition pinning."""
        return self._record(
            AuditEventType.PIN_CREATED,
            tool_name,
            details={"definition_hash": definition_hash},
        )

    def log_pin_mismatch(
        self,
        tool_name: str,
        expected_hash: str,
        actual_hash: str,
    ) -> AuditEvent:
        """Record a tool definition mismatch (potential rug pull)."""
        return self._record(
            AuditEventType.PIN_MISMATCH,
            tool_name,
            details={
                "expected_hash": expected_hash,
                "actual_hash": actual_hash,
            },
        )

    def log_replay_detected(
        self,
        tool_name: str,
        nonce: str,
    ) -> AuditEvent:
        """Record a replay attack detection."""
        return self._record(
            AuditEventType.REPLAY_DETECTED,
            tool_name,
            nonce=nonce,
            details={"alert": "Duplicate nonce detected"},
        )

    def get_events_for_tool(self, tool_name: str) -> List[AuditEvent]:
        """Get all events for a specific tool."""
        with self._lock:
            return [e for e in self._events if e.tool_name == tool_name]

    def get_events_by_type(self, event_type: str) -> List[AuditEvent]:
        """Get all events of a given type."""
        with self._lock:
            return [e for e in self._events if e.event_type == event_type]

    def get_errors(self) -> List[AuditEvent]:
        """Get all error events."""
        return self.get_events_by_type(AuditEventType.TOOL_ERROR)

    def get_security_alerts(self) -> List[AuditEvent]:
        """Get all security-related events (mismatches, replays)."""
        alert_types = {
            AuditEventType.PIN_MISMATCH,
            AuditEventType.REPLAY_DETECTED,
        }
        with self._lock:
            return [e for e in self._events if e.event_type in alert_types]

    def summary(self) -> Dict[str, Any]:
        """Generate a summary of the audit trail."""
        with self._lock:
            events = list(self._events)

        tool_calls = [e for e in events if e.event_type == AuditEventType.TOOL_CALL]
        errors = [e for e in events if e.event_type == AuditEventType.TOOL_ERROR]
        alerts = [
            e for e in events
            if e.event_type in {AuditEventType.PIN_MISMATCH, AuditEventType.REPLAY_DETECTED}
        ]

        tools_called = set()
        for e in tool_calls:
            tools_called.add(e.tool_name)

        return {
            "total_events": len(events),
            "total_calls": len(tool_calls),
            "total_errors": len(errors),
            "security_alerts": len(alerts),
            "unique_tools": len(tools_called),
            "tools": list(tools_called),
            "first_event": events[0].timestamp if events else None,
            "last_event": events[-1].timestamp if events else None,
        }

    def export_json(self, filepath: str) -> None:
        """Export the full audit trail to a JSON file."""
        with self._lock:
            events_snapshot = list(self._events)

        # Build summary outside the lock to avoid deadlock
        summary = self.summary()
        data = {
            "audit_trail": [e.to_dict() for e in events_snapshot],
            "summary": summary,
            "exported_at": time.time(),
        }
        with open(filepath, "w") as f:
            json.dump(data, f, indent=2)

    def clear(self) -> None:
        """Clear all events. Use with caution."""
        with self._lock:
            self._events.clear()
            self._sequence = 0

    def _record(
        self,
        event_type: str,
        tool_name: str,
        nonce: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
    ) -> AuditEvent:
        """Record an event to the trail."""
        with self._lock:
            self._sequence += 1
            event = AuditEvent(
                event_type=event_type,
                tool_name=tool_name,
                timestamp=time.time(),
                nonce=nonce,
                details=details or {},
                sequence=self._sequence,
            )
            self._events.append(event)

            # Trim if over max
            if len(self._events) > self._max_events:
                self._events = self._events[-self._max_events:]

        # Notify callback if registered
        if self._on_event:
            try:
                self._on_event(event)
            except Exception:
                pass  # Don't let callback errors break the audit

        return event
