"""Hash chain implementation for immutable event ordering.

Each agent maintains an independent hash chain. Every significant event
(baseline creation, drift test, key rotation) is appended to the chain.

Chain hash formula:
    chain_hash = SHA-256(content_hash || prev_chain_hash)

The chain provides:
    - Tamper detection (any modification breaks the chain)
    - Temporal ordering (sequence numbers are strictly monotonic)
    - Independent verification (anyone with the chain can verify integrity)
"""

from __future__ import annotations

import hashlib
from datetime import datetime, timezone
from typing import Optional, Protocol

from pydantic import BaseModel, Field


class ChainEntryModel(BaseModel):
    """A single entry in an agent's hash chain."""
    agent_pubkey: str
    sequence: int
    event_type: str  # "baseline", "test", "key_rotation", "recovery"
    content_hash: str  # SHA-256 of the referenced object
    prev_hash: Optional[str] = None  # previous chain_hash (None for genesis)
    chain_hash: str = ""  # SHA-256(content_hash || prev_hash)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class ChainStore(Protocol):
    """Protocol for chain storage backends."""

    def get_latest_entry(self, agent_pubkey: str) -> Optional[ChainEntryModel]:
        """Get the latest chain entry for an agent."""
        ...

    def get_chain(self, agent_pubkey: str) -> list[ChainEntryModel]:
        """Get all chain entries for an agent, ordered by sequence."""
        ...

    def save_entry(self, entry: ChainEntryModel) -> None:
        """Save a chain entry."""
        ...


def compute_chain_hash(content_hash: str, prev_hash: Optional[str]) -> str:
    """Compute the chain hash: SHA-256(content_hash || prev_hash).

    For the genesis entry (no previous), prev_hash is treated as empty string.
    """
    # WHY empty string for genesis: Using "" (not a zero hash) for the genesis
    # entry's prev_hash makes verification simpler -- we don't need a special
    # case for "is this a zero hash or a real hash?" The chain formula is
    # unambiguous: SHA-256(content || "") for genesis, SHA-256(content || prev)
    # for all subsequent entries.
    prev = prev_hash or ""
    payload = (content_hash + prev).encode("utf-8")
    return hashlib.sha256(payload).hexdigest()


def append_to_chain(
    agent_pubkey: str,
    event_type: str,
    content_hash: str,
    store: ChainStore,
) -> ChainEntryModel:
    """Append a new entry to an agent's hash chain.

    Args:
        agent_pubkey: The agent's public key.
        event_type: Type of event ("baseline", "test", "key_rotation", "recovery").
        content_hash: SHA-256 hash of the content being recorded.
        store: Storage backend implementing ChainStore protocol.

    Returns:
        The new ChainEntryModel with computed chain_hash.
    """
    latest = store.get_latest_entry(agent_pubkey)

    if latest is None:
        sequence = 0
        prev_hash = None
    else:
        sequence = latest.sequence + 1
        prev_hash = latest.chain_hash

    chain_hash = compute_chain_hash(content_hash, prev_hash)

    entry = ChainEntryModel(
        agent_pubkey=agent_pubkey,
        sequence=sequence,
        event_type=event_type,
        content_hash=content_hash,
        prev_hash=prev_hash,
        chain_hash=chain_hash,
    )

    store.save_entry(entry)
    return entry


def verify_chain(
    agent_pubkey: str,
    store: ChainStore,
) -> tuple[bool, Optional[int]]:
    """Verify the integrity of an agent's hash chain.

    Args:
        agent_pubkey: The agent's public key.
        store: Storage backend implementing ChainStore protocol.

    Returns:
        (is_valid, break_point) — True if chain is intact, or
        False with the sequence number where the break was detected.
    """
    entries = store.get_chain(agent_pubkey)

    if not entries:
        return True, None

    # Verify genesis entry
    expected_hash = compute_chain_hash(entries[0].content_hash, None)
    if entries[0].chain_hash != expected_hash:
        return False, 0
    if entries[0].prev_hash is not None:
        return False, 0

    # Verify each subsequent entry
    for i in range(1, len(entries)):
        entry = entries[i]
        prev_entry = entries[i - 1]

        # Sequence must be strictly monotonic
        if entry.sequence != prev_entry.sequence + 1:
            return False, entry.sequence

        # prev_hash must reference the previous entry's chain_hash
        if entry.prev_hash != prev_entry.chain_hash:
            return False, entry.sequence

        # Recompute and verify chain_hash
        expected = compute_chain_hash(entry.content_hash, entry.prev_hash)
        if entry.chain_hash != expected:
            return False, entry.sequence

    return True, None
