"""Canonical encoding for memory snapshots and branch summaries.

Determinism rules (per ticket #000017):

- Branch IDs are stable strings (``failure-motif:title-mismatch``,
  ``correction:operator-applied``). Adding a branch type is a new
  ``schema_version`` (``memory-v1`` → ``memory-v2``).
- ``branch_digest`` over the (branch_id, summary_blob, count) triple
  pins each branch independently — re-running the projection on
  the same audit slice produces the same digest.
- ``memory_root`` is SHA-256 over the canonical-JSON of:
  ``{schema_version, parent_memory_root, audit_events_high_water,
  branch_digests}`` where ``branch_digests`` is a sorted list of
  ``{"branch_id": ..., "digest": ..., "count": ...}`` triples.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from typing import Optional


SCHEMA_VERSION = "memory-v1"


@dataclass(frozen=True)
class BranchSummary:
    """One branch of a memory snapshot.

    ``summary_blob`` is the deterministic projection output for this
    branch — its content is opaque to ``canonical.py`` but must be
    bytes-stable for the same audit slice.
    """

    branch_id: str
    summary_blob: bytes
    count: int

    def canonical(self) -> dict:
        return {
            "branch_id": self.branch_id,
            "digest": branch_digest(self),
            "count": int(self.count),
        }


def branch_digest(summary: BranchSummary) -> str:
    """SHA-256 over (branch_id, summary_blob bytes, count)."""
    h = hashlib.sha256()
    h.update(summary.branch_id.encode("utf-8", errors="surrogatepass"))
    h.update(b"\x00")
    h.update(summary.summary_blob)
    h.update(b"\x00")
    h.update(str(int(summary.count)).encode("ascii"))
    return h.hexdigest()


@dataclass(frozen=True)
class MemorySnapshot:
    """Memory snapshot record body (pre-persistence).

    ``audit_events_high_water`` is the ``event_hash`` of the most
    recent audit event included in the projection — this lets a
    later snapshot resume from where this one stopped without
    re-summarizing the whole chain.
    """

    schema_version: str
    parent_memory_root: Optional[str]
    audit_events_high_water: str
    branches: tuple[BranchSummary, ...] = field(default_factory=tuple)

    def canonical(self) -> dict:
        # branches are sorted by branch_id so insertion order doesn't
        # affect the memory_root.
        sorted_branches = sorted(self.branches, key=lambda b: b.branch_id)
        return {
            "schema_version": self.schema_version,
            "parent_memory_root": self.parent_memory_root,
            "audit_events_high_water": self.audit_events_high_water,
            "branches": [b.canonical() for b in sorted_branches],
        }


def canonical_branches_blob(snapshot: MemorySnapshot) -> str:
    return json.dumps(
        snapshot.canonical(),
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )


def memory_root(snapshot: MemorySnapshot) -> str:
    """SHA-256 over the canonical body bytes."""
    return hashlib.sha256(
        canonical_branches_blob(snapshot).encode(
            "utf-8", errors="surrogatepass"
        )
    ).hexdigest()
