"""Bounded exhaustive model comparing two NIP-RS manual-unread encodings.

Candidate A: lexicographic operation register
  Per context: {counter, client_tiebreak, op in {SET,CLEAR}, baseline}
  in a NEW top-level field beside `contexts`.
  Merge = max tuple (counter, tiebreak, op-rule on full tie).

Candidate B: two grow-only counters + baseline
  Per context: S (set counter), C (clear counter), B (frontier-at-set-time)
  as sibling keys under `contexts` (ov_s:, ov_c:, ov_b: prefixes).
  Action: own counter := max(S,C)+1; set also writes B := effective frontier.
  Merge = componentwise max. Tie policy on S == C is a parameter.

Both share:
  - Frontier: grow-only max() per NIP-RS v1 (unchanged).
  - Verdict: unread(ctx) = latest > effective_frontier(ctx) OR override_set(ctx).
  - Mark-read = advance frontier + clear override.
  - Mark-unread = set override with baseline B = current effective frontier.
  - Natural frontier advance strictly past B dominates a stale set.

Device simulators use overridable methods (_override_set, _compact, _merge_reg,
_bump, _sanitize_value) so the mutation harness can inject weakened rules via
subclassing without monkeypatching.
"""
from dataclasses import dataclass
from typing import Optional


SET = "SET"
CLEAR = "CLEAR"


# ---------------------------------------------------------------------------
# Reserved key namespace + escaping
#
# NIP-RS v1 context IDs are arbitrary UTF-8 (spec :89, :113-114), so a
# pre-existing opaque context could legitimately begin with `ov_s:`,
# `ov_c:`, or `ov_b:` and collide with a control key for a DIFFERENT
# context in the same flattened `contexts` map. `ov_` (the shared
# 3-byte stem) and the escape marker itself are reserved; any raw
# context ID that would collide is escaped before being used as a
# plain frontier key. Escaping is a no-op for every context ID Buzz
# actually generates (channel UUID, `msg:<hex64>`, `thread:<hex64>`
# — none start with `ov_` or `esc:`), so the common case pays zero
# bytes. Only a pathological ID pays the 4-byte `esc:` cost.
#
# This protects context IDs generated by amendment-aware clients.
# It does NOT retroactively protect a context that a PRE-EXISTING
# legacy client already published unescaped before the amendment
# shipped — that residual hazard is documented, not solved (see
# NOTE.md "Reserved key namespace").
# ---------------------------------------------------------------------------

ESCAPE_PREFIX = "esc:"
_RESERVED_STEM = "ov_"


def _needs_escape(raw_key: str) -> bool:
    return raw_key.startswith(_RESERVED_STEM) or raw_key.startswith(ESCAPE_PREFIX)


def escape_context_key(raw_key: str) -> str:
    return ESCAPE_PREFIX + raw_key if _needs_escape(raw_key) else raw_key


def unescape_context_key(wire_key: str) -> str:
    if wire_key.startswith(ESCAPE_PREFIX):
        return wire_key[len(ESCAPE_PREFIX):]
    return wire_key


# ---------------------------------------------------------------------------
# Candidate B — two grow-only counters + baseline
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class RegB:
    s: int = 0
    c: int = 0
    b: int = 0


def merge_reg_b(a: Optional[RegB], b: Optional[RegB]) -> Optional[RegB]:
    if a is None:
        return b
    if b is None:
        return a
    return RegB(s=max(a.s, b.s), c=max(a.c, b.c), b=max(a.b, b.b))


def override_set_b(reg: Optional[RegB], frontier_val: int, tie_policy=CLEAR) -> bool:
    if reg is None:
        return False
    if frontier_val > reg.b and reg.s > 0:
        return False
    if reg.s > reg.c:
        return True
    if reg.s == reg.c and reg.s > 0:
        return tie_policy == SET
    return False


def compact_b(reg: RegB, frontier_val: int, tie_policy=CLEAR) -> Optional[RegB]:
    """Compact override state.

    Tombstone-floor design: a register with any recorded counter
    activity (S>0 or C>0) is never fully deleted. Its counter
    high-water-mark is exactly what prevents a stale replica —
    any (S,C) pair below that ceiling — from dominating a freshly
    created register after compaction (delete-on-dominance made
    counters reusable: a dead register dropped entirely, then a new
    local set/clear pair restarted from S=0/C=0, so a delayed stale
    peer snapshot with S>0 could out-rank the new state on replay).
    Only a virgin register (S==0, C==0, no activity ever recorded)
    has no ceiling to protect and compacts to None.

    A live override (per `override_set_b`, which is already
    policy-aware) is returned unchanged — compaction only touches dead
    state. Dead overrides — whether dominated by C>S, tied under
    clear-wins, or baseline-dominated by frontier advance — compact to
    the clear-tombstone floor `RegB(s=0, c=max(S,C), b=0)`: S is
    zeroed (no longer overriding), but C retains the ceiling so both a
    future local bump (`max(S,C)+1`) and a componentwise-max merge with
    any pre-compaction stale snapshot start strictly above the
    historical maximum, never below it.
    """
    if reg.s == 0 and reg.c == 0:
        return None
    if override_set_b(reg, frontier_val, tie_policy):
        return reg
    return RegB(s=0, c=max(reg.s, reg.c), b=0)


# ---------------------------------------------------------------------------
# Candidate A — lexicographic operation register
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class RegA:
    counter: int = 0
    tiebreak: str = ""
    op: str = CLEAR
    baseline: int = 0

    def as_tuple(self, op_wins):
        op_val = 1 if self.op == op_wins else 0
        return (self.counter, self.tiebreak, op_val)


def merge_reg_a(a: Optional[RegA], b: Optional[RegA], tie_op=CLEAR) -> Optional[RegA]:
    if a is None:
        return b
    if b is None:
        return a
    at = a.as_tuple(tie_op)
    bt = b.as_tuple(tie_op)
    if at == bt:
        return RegA(
            counter=a.counter, tiebreak=a.tiebreak, op=a.op,
            baseline=max(a.baseline, b.baseline),
        )
    return a if at > bt else b


# ---------------------------------------------------------------------------
# Device simulation — Candidate B
# ---------------------------------------------------------------------------

class DeviceB:
    """Simulates one device's NIP-RS read-state blob with manual-unread
    override layer (candidate B encoding).

    All model operations go through overridable _methods so the mutation
    harness can inject weakened rules via subclassing.
    """

    def __init__(self, client_id, is_legacy=False):
        self.client_id = client_id
        self.is_legacy = is_legacy
        self.frontier = {}
        self.overrides = {}

    def effective_frontier(self, ctx):
        return self.frontier.get(ctx, 0)

    def _override_set(self, reg, frontier_val, tie_policy):
        return override_set_b(reg, frontier_val, tie_policy)

    def _compact(self, reg, frontier_val, tie_policy):
        return compact_b(reg, frontier_val, tie_policy)

    def _merge_reg(self, a, b):
        return merge_reg_b(a, b)

    def _bump(self, s, c):
        return max(s, c) + 1

    def _sanitize_value(self, v):
        return isinstance(v, int) and 0 <= v <= 4294967295

    def override_is_set(self, ctx, tie_policy=CLEAR):
        return self._override_set(
            self.overrides.get(ctx), self.effective_frontier(ctx), tie_policy
        )

    def verdict(self, ctx, latest_ts, tie_policy=CLEAR):
        return (latest_ts > self.effective_frontier(ctx)
                or self.override_is_set(ctx, tie_policy))

    def do_mark_unread(self, ctx):
        if self.is_legacy:
            return
        cur = self.overrides.get(ctx, RegB())
        new_s = self._bump(cur.s, cur.c)
        self.overrides[ctx] = RegB(s=new_s, c=cur.c, b=self.effective_frontier(ctx))

    def do_mark_read(self, ctx, frontier_ts):
        self.frontier[ctx] = max(self.frontier.get(ctx, 0), frontier_ts)
        if not self.is_legacy:
            cur = self.overrides.get(ctx, RegB())
            new_c = self._bump(cur.s, cur.c)
            self.overrides[ctx] = RegB(s=cur.s, c=new_c, b=cur.b)

    def do_advance_frontier(self, ctx, ts):
        self.frontier[ctx] = max(self.frontier.get(ctx, 0), ts)

    def do_compact(self, ctx, tie_policy=CLEAR):
        reg = self.overrides.get(ctx)
        if reg is None:
            return
        result = self._compact(reg, self.effective_frontier(ctx), tie_policy)
        if result is None:
            if ctx in self.overrides:
                del self.overrides[ctx]
        else:
            self.overrides[ctx] = result

    def do_reinstall(self):
        self.client_id = self.client_id + "_r"
        self.frontier = {}
        self.overrides = {}

    def _canonicalize_for_publish(self, ctx, tie_policy):
        """Canonical published form of `ctx`'s override register,
        computed fresh against the current effective frontier —
        independent of whether `do_compact` was ever called locally.
        Returns `(is_live, canonical_reg)`; `canonical_reg is None`
        means virgin (omit from the wire entirely). Reuses the same
        overridable `_compact`/`_override_set` hooks `do_compact` uses,
        so a mutation-harness subclass that weakens one weakens both
        the storage-GC path and the publish path identically.
        """
        reg = self.overrides.get(ctx)
        if reg is None:
            return False, None
        front = self.effective_frontier(ctx)
        canonical = self._compact(reg, front, tie_policy)
        if canonical is None:
            return False, None
        return self._override_set(canonical, front, tie_policy), canonical

    def publish_blob(self, tie_policy=CLEAR):
        """Serialize this device's read-state blob.

        Every override is canonicalized at serialization time: live ->
        unchanged (3 keys), dead -> tombstone floor (1 key, `ov_c:`
        only), virgin -> omitted (0 keys). Canonical publication is a
        protocol requirement, not an optimization — noncanonical wire
        output is structurally impossible here, not merely avoided by
        convention. `do_compact` remains a separate storage-GC
        transition that mutates `self.overrides`; publication no
        longer depends on it having been called first.

        **Atomic slot-grouping rule (spec-amendment requirement):**
        A context's frontier entry and ALL of its `ov_*` sibling entries
        MUST travel in the same slot. `split_blob_into_slots` below
        enforces this by round-robining per-context groups, never
        per-entry. A receiving client that only holds part of a context
        group and attempts to reconstruct a `RegB` from it would see
        partial zeroes and might canonically re-publish a false
        tombstone. Group atomicity makes partial reconstruction
        structurally impossible from a compliant publisher's output.
        """
        blob_ctx = {escape_context_key(k): v for k, v in self.frontier.items()}
        if not self.is_legacy:
            for k in self.overrides:
                is_live, canonical = self._canonicalize_for_publish(k, tie_policy)
                if canonical is None:
                    continue  # virgin: omitted from the wire entirely
                if is_live:
                    blob_ctx[f"ov_s:{k}"] = canonical.s
                    blob_ctx[f"ov_c:{k}"] = canonical.c
                    blob_ctx[f"ov_b:{k}"] = canonical.b
                else:
                    blob_ctx[f"ov_c:{k}"] = canonical.c  # tombstone: ceiling only
        return {"v": 1, "client_id": self.client_id, "contexts": blob_ctx}

    def split_blob_into_slots(self, tie_policy=CLEAR, n_slots=2):
        """Split this device's blob into `n_slots` compliant slots.

        **Atomic grouping rule:** a context's frontier entry and ALL of
        its `ov_*` sibling entries travel together in the same slot.
        Round-robin assignment is per-context group, never per-entry.
        This matches production `splitContextsIntoBudgetedSlots` when
        it is amended to group by context instead of by individual entry.

        Returns a list of `n_slots` blobs, each with the same `v` and
        `client_id` but a disjoint subset of context groups.
        """
        blob = self.publish_blob(tie_policy)
        contexts = blob["contexts"]

        # Gather per-context groups: each group is a list of (key, value) pairs.
        # A "group" is: the frontier key (escaped ctx) + any ov_* siblings.
        # Contexts that appear only as ov_* keys (no frontier entry) are
        # also grouped together.
        groups = {}  # logical_ctx -> list of (wire_key, value)
        for wire_key, value in contexts.items():
            if wire_key.startswith("ov_s:"):
                ctx = wire_key[5:]
            elif wire_key.startswith("ov_c:"):
                ctx = wire_key[5:]
            elif wire_key.startswith("ov_b:"):
                ctx = wire_key[5:]
            else:
                # Frontier key: may be escaped (e.g. "esc:ov_s:evil").
                # Derive the logical context ID by unescaping so this
                # entry joins the same group as its ov_* siblings, which
                # are keyed by the RAW context ID (e.g. "ov_s:evil" ->
                # ctx = "evil", but "esc:ov_s:evil" frontier -> ctx =
                # "ov_s:evil" after unescape).  Without this step an
                # escaped frontier key and its ov_* siblings would be
                # treated as two different groups, splitting the register
                # across slots — reproducing the round-1 partial-
                # reconstruction poison for escaped context IDs.
                ctx = unescape_context_key(wire_key)
            groups.setdefault(ctx, []).append((wire_key, value))

        slots = [{"v": blob["v"], "client_id": blob["client_id"], "contexts": {}}
                 for _ in range(n_slots)]
        for i, (_ctx, pairs) in enumerate(sorted(groups.items())):
            slot = slots[i % n_slots]
            for wire_key, value in pairs:
                slot["contexts"][wire_key] = value
        return slots

    def receive_merge(self, blob):
        incoming_overrides = {}
        for k, v in blob.get("contexts", {}).items():
            if k.startswith("ov_s:"):
                ctx = k[5:]
                incoming_overrides.setdefault(ctx, [0, 0, 0])[0] = v
            elif k.startswith("ov_c:"):
                ctx = k[5:]
                incoming_overrides.setdefault(ctx, [0, 0, 0])[1] = v
            elif k.startswith("ov_b:"):
                ctx = k[5:]
                incoming_overrides.setdefault(ctx, [0, 0, 0])[2] = v
            else:
                ctx = unescape_context_key(k)
                self.frontier[ctx] = max(self.frontier.get(ctx, 0), v)

        if not self.is_legacy:
            for ctx, (s, c, b) in incoming_overrides.items():
                incoming_reg = RegB(s=s, c=c, b=b)
                self.overrides[ctx] = self._merge_reg(
                    self.overrides.get(ctx), incoming_reg
                )

    def legacy_sanitize_and_publish(self, tie_policy=CLEAR):
        blob = self.publish_blob(tie_policy)
        sanitized = {}
        for k, v in blob["contexts"].items():
            if len(k.encode("utf-8")) <= 256 and self._sanitize_value(v):
                sanitized[k] = v
        return {"v": 1, "client_id": self.client_id, "contexts": sanitized}

    def state_key(self, contexts, tie_policy=CLEAR):
        parts = []
        for ctx in sorted(contexts):
            f = self.effective_frontier(ctx)
            reg = self.overrides.get(ctx, RegB())
            ov = self.override_is_set(ctx, tie_policy)
            parts.append((ctx, f, reg.s, reg.c, reg.b, ov))
        return (self.client_id, self.is_legacy, tuple(parts))


# ---------------------------------------------------------------------------
# Device simulation — Candidate A
# ---------------------------------------------------------------------------

class DeviceA:
    def __init__(self, client_id, is_legacy=False):
        self.client_id = client_id
        self.is_legacy = is_legacy
        self.frontier = {}
        self.overrides = {}
        self.counter = 0

    def effective_frontier(self, ctx):
        return self.frontier.get(ctx, 0)

    def override_is_set(self, ctx):
        reg = self.overrides.get(ctx)
        if reg is None or reg.op == CLEAR:
            return False
        if self.effective_frontier(ctx) > reg.baseline:
            return False
        return True

    def verdict(self, ctx, latest_ts):
        return latest_ts > self.effective_frontier(ctx) or self.override_is_set(ctx)

    def do_mark_unread(self, ctx):
        if self.is_legacy:
            return
        self.counter += 1
        self.overrides[ctx] = RegA(
            counter=self.counter, tiebreak=self.client_id,
            op=SET, baseline=self.effective_frontier(ctx),
        )

    def do_mark_read(self, ctx, frontier_ts):
        self.frontier[ctx] = max(self.frontier.get(ctx, 0), frontier_ts)
        if not self.is_legacy:
            self.counter += 1
            self.overrides[ctx] = RegA(
                counter=self.counter, tiebreak=self.client_id,
                op=CLEAR, baseline=0,
            )

    def do_advance_frontier(self, ctx, ts):
        self.frontier[ctx] = max(self.frontier.get(ctx, 0), ts)

    def receive_merge(self, blob, tie_op=CLEAR):
        for ctx, ts in blob.get("contexts", {}).items():
            self.frontier[ctx] = max(self.frontier.get(ctx, 0), ts)
        if not self.is_legacy:
            for ctx, reg in blob.get("overrides", {}).items():
                self.overrides[ctx] = merge_reg_a(
                    self.overrides.get(ctx), reg, tie_op
                )
                if reg.counter > self.counter:
                    self.counter = reg.counter

    def publish_blob(self):
        blob = {"v": 1, "client_id": self.client_id, "contexts": dict(self.frontier)}
        if not self.is_legacy:
            blob["overrides"] = dict(self.overrides)
        return blob

    def legacy_rewrite_and_publish(self):
        return {"v": 1, "client_id": self.client_id, "contexts": dict(self.frontier)}


# ---------------------------------------------------------------------------
# Legacy pruning/trim model
# ---------------------------------------------------------------------------

def legacy_prune(contexts, horizon):
    return {k: v for k, v in contexts.items()
            if not (k.startswith("msg:") or k.startswith("thread:")) or v >= horizon}


def legacy_trim(contexts, client_id, max_bytes=32768):
    import json

    def size(ctx):
        return len(json.dumps({"v": 1, "client_id": client_id, "contexts": ctx}).encode())

    if size(contexts) <= max_bytes:
        return contexts, True
    evictable = sorted(
        ((k, v) for k, v in contexts.items()
         if k.startswith("msg:") or k.startswith("thread:")),
        key=lambda kv: kv[1],
    )
    out = dict(contexts)
    for k, _ in evictable:
        del out[k]
        if size(out) <= max_bytes:
            return out, True
    return out, size(out) <= max_bytes


def legacy_sanitize_blob(blob):
    sanitized = {}
    for k, v in blob.get("contexts", {}).items():
        if (len(k.encode("utf-8")) <= 256
                and isinstance(v, int) and 0 <= v <= 4294967295):
            sanitized[k] = v
    return {"v": 1, "client_id": blob.get("client_id", ""), "contexts": sanitized}