"""
Merkle Mountain Range (MMR) audit chain.

Provides tamper-evident sequential logging where each entry commits to
all previous entries via a hash chain. Any tampering with a historical
entry invalidates all subsequent entries.

This is a simplified MMR (append-only, no peak merging for simplicity)
that provides the core guarantee: each leaf_hash includes the previous
root, making the chain tamper-evident.

Usage::

    chain = AuditChain()
    entry = chain.append(action_type="tool_executed", payload={"tool": "web_search"})
    # entry.leaf_hash, entry.mmr_root, entry.leaf_index

    # Or use the module-level singleton:
    from merkle_audit import log_event
    entry = log_event("tool_executed", {"tool": "web_search", "query": "..."})
"""
from __future__ import annotations

import hashlib
import json
import threading
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional


@dataclass
class ChainEntry:
    leaf_index: int
    leaf_hash: str       # SHA-256 of (prev_root + action_type + payload_hash + timestamp)
    mmr_root: str        # SHA-256 of all leaf hashes up to this point (running root)
    payload_hash: str    # SHA-256 of the serialized payload
    action_type: str
    actor: str           # Who performed the action (agent ID, user, service name, etc.)
    agent_id: str        # Deprecated alias for actor — kept for compatibility
    created_at: float


class AuditChain:
    """
    Thread-safe in-memory MMR audit chain with persistence hooks.

    The chain guarantees: if any historical entry is modified,
    all subsequent leaf_hash values will not match their expected values.

    For persistence, register a ``persist_hook`` — a callable that accepts
    a :class:`ChainEntry` and stores it (SQLite, Postgres, file, etc.).
    Persistence failures are silently swallowed so the chain never crashes
    your application.

    Example::

        chain = AuditChain()

        # Attach SQLite persistence (from merkle_audit.persister):
        from merkle_audit.persister import SQLitePersister
        chain.set_persist_hook(SQLitePersister("audit.db"))

        entry = chain.append("payment_processed", {"amount": 42.00, "currency": "USD"})
        print(entry.leaf_hash, entry.mmr_root)
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._leaf_count: int = 0
        self._current_root: str = "0" * 64  # Genesis root (64 zero hex chars)
        self._persist_hook: Optional[Callable[[ChainEntry], None]] = None

    def set_persist_hook(self, hook: Callable[[ChainEntry], None]) -> None:
        """Register a callable hook(entry: ChainEntry) for persistence."""
        self._persist_hook = hook

    def append(
        self,
        action_type: str,
        payload: Dict[str, Any],
        actor: str = "",
        agent_id: str = "",
    ) -> ChainEntry:
        """
        Append a new entry to the chain. Thread-safe.

        Args:
            action_type: A string identifying the action (e.g. ``"tool_executed"``).
            payload: Arbitrary JSON-serializable dict describing the event.
            actor: Who performed the action (agent ID, username, service name, etc.).
            agent_id: Alias for ``actor`` (legacy parameter name).

        Returns:
            The :class:`ChainEntry` with all computed hashes.
        """
        with self._lock:
            now = time.time()
            payload_str = json.dumps(payload, sort_keys=True, default=str)
            payload_hash = hashlib.sha256(payload_str.encode()).hexdigest()

            # leaf_hash commits to: previous root + action + payload + timestamp
            leaf_content = f"{self._current_root}:{action_type}:{payload_hash}:{now}"
            leaf_hash = hashlib.sha256(leaf_content.encode()).hexdigest()

            # New root commits to: previous root + new leaf
            new_root = hashlib.sha256(
                f"{self._current_root}:{leaf_hash}".encode()
            ).hexdigest()

            resolved_actor = actor or agent_id
            entry = ChainEntry(
                leaf_index=self._leaf_count,
                leaf_hash=leaf_hash,
                mmr_root=new_root,
                payload_hash=payload_hash,
                action_type=action_type,
                actor=resolved_actor,
                agent_id=resolved_actor,
                created_at=now,
            )

            self._leaf_count += 1
            self._current_root = new_root
            hook = self._persist_hook

        # Call the persist hook outside the lock to avoid deadlocks
        if hook is not None:
            try:
                hook(entry)
            except Exception:
                pass  # Persistence failures must never crash the chain

        return entry

    def current_root(self) -> str:
        """Return the current MMR root (hex string)."""
        with self._lock:
            return self._current_root

    def leaf_count(self) -> int:
        """Return the current number of entries."""
        with self._lock:
            return self._leaf_count

    def verify_entry(self, entry: ChainEntry, prev_root: str) -> bool:
        """
        Verify that an entry is consistent with a known previous root.

        Use this to detect tampering: if any historical entry was modified,
        recomputing its ``leaf_hash`` from ``prev_root`` will produce a
        different value than what was stored.

        Args:
            entry: The :class:`ChainEntry` to verify.
            prev_root: The MMR root immediately *before* this entry was appended.
                       For the first entry, use ``"0" * 64``.

        Returns:
            ``True`` if the entry is valid, ``False`` if tampered.
        """
        expected_leaf = (
            f"{prev_root}:{entry.action_type}:{entry.payload_hash}:{entry.created_at}"
        )
        expected_leaf_hash = hashlib.sha256(expected_leaf.encode()).hexdigest()
        if entry.leaf_hash != expected_leaf_hash:
            return False
        expected_root = hashlib.sha256(
            f"{prev_root}:{entry.leaf_hash}".encode()
        ).hexdigest()
        return entry.mmr_root == expected_root


# ---------------------------------------------------------------------------
# Module-level singleton
# ---------------------------------------------------------------------------

_global_chain: Optional[AuditChain] = None


def get_chain() -> AuditChain:
    """Get or create the global audit chain singleton."""
    global _global_chain
    if _global_chain is None:
        _global_chain = AuditChain()
    return _global_chain


def log_event(
    action_type: str,
    payload: Dict[str, Any],
    actor: str = "",
    agent_id: str = "",
) -> ChainEntry:
    """
    Append an event to the global audit chain.

    This is the primary convenience API. For production use with persistence,
    call ``get_chain().set_persist_hook(...)`` once at startup.

    Example::

        from merkle_audit import log_event

        entry = log_event(
            "tool_executed",
            {"tool": "web_search", "query": "climate change"},
            actor="research-agent-1",
        )
    """
    return get_chain().append(
        action_type=action_type,
        payload=payload,
        actor=actor,
        agent_id=agent_id,
    )
