"""CLERK: Consolidated Ledger with Eviction and Rewrite Keys.

Core memory machinery shared by the data generator, the training harness and
the evaluators:

  * ``Slot``        -- one memory slot (a typed fact plus lifecycle metadata)
  * edit-op schema  -- ADD / UPDATE / TOMBSTONE / EVICT / NOOP
  * ``apply_ops``   -- the DETERMINISTIC reducer that applies a model-emitted
                       edit program to a ledger
  * serializers     -- compact JSON forms of ledgers and op programs

The reducer is the single source of truth: the generator emits gold op
programs, and every downstream consumer applies them through ``apply_ops``.
This guarantees that the training supervision, the simulated model rollouts
and the evaluation all execute the exact same state transition.

Author: Justin Wolcott (fallnai-research.org)
License: Apache-2.0
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

LEDGER_BUDGET_DEFAULT = 24
VALID_OPS = ("ADD", "UPDATE", "TOMBSTONE", "EVICT", "NOOP")


@dataclass
class Slot:
    """A single ledger slot."""

    id: int
    subject: str = ""
    predicate: str = ""
    object: str = ""
    t_last: int = -1          # session index when this slot was last written
    n_uses: int = 0           # number of times the fact has been queried (internal)
    valid: bool = True        # False => tombstoned (fact no longer true)
    occupied: bool = False    # False => free (empty, or freed by EVICT)


def empty_ledger(budget: int = LEDGER_BUDGET_DEFAULT) -> List[Slot]:
    return [Slot(id=i) for i in range(budget)]


def serialize_ledger(ledger: List[Slot]) -> str:
    """Compact JSON view of a ledger, as seen by the model.

    Tombstoned slots stay visible (with v=0) until they are evicted, so
    temporal questions about superseded values remain answerable.
    """
    rows = []
    for s in ledger:
        if s.occupied:
            rows.append(
                {"i": s.id, "s": s.subject, "p": s.predicate,
                 "o": s.object, "t": s.t_last, "v": 1 if s.valid else 0}
            )
    return json.dumps(rows, separators=(",", ":"), ensure_ascii=False)


def parse_ledger(text: str, budget: int = LEDGER_BUDGET_DEFAULT) -> List[Slot]:
    """Rebuild a ledger from its serialized form."""
    ledger = empty_ledger(budget)
    for row in json.loads(text):
        s = ledger[row["i"]]
        s.occupied = True
        s.subject = row["s"]
        s.predicate = row["p"]
        s.object = row["o"]
        s.t_last = row["t"]
        s.valid = bool(row["v"])
    return ledger


def parse_ops(text: str) -> Tuple[List[Dict[str, Any]], str]:
    """Parse a model-emitted edit program.

    Returns (ops, error). ``error`` is "" on success, otherwise a short
    description of the first problem found. Malformed programs are surfaced
    to the evaluation harness instead of being silently dropped.
    """
    text = text.strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.startswith("json"):
            text = text[4:]
    try:
        payload = json.loads(text)
    except json.JSONDecodeError as e:
        return [], f"json_decode_error: {e}"
    if not isinstance(payload, dict) or "ops" not in payload:
        return [], "missing_top_level_ops_key"
    ops = payload["ops"]
    if not isinstance(ops, list):
        return [], "ops_not_a_list"
    clean: List[Dict[str, Any]] = []
    for op in ops:
        if not isinstance(op, dict) or "op" not in op:
            return clean, "op_missing_op_key"
        if op["op"] not in VALID_OPS:
            return clean, f"unknown_op:{op['op']}"
        clean.append(op)
    return clean, ""


def apply_ops(
    ledger: List[Slot],
    ops: List[Dict[str, Any]],
    session_idx: int,
) -> Tuple[List[Slot], Dict[str, int]]:
    """Apply an edit program to a ledger. Deterministic; mutates a copy.

    Rules:
      ADD       -- write into the lowest free slot; if the ledger has no free
                   slot the op is DROPPED and counted (a trained policy must
                   EVICT first; the gold generator always does).
      UPDATE    -- overwrite s/p/o of slot ``i`` and bump t_last. The slot
                   keeps its identity, so supersession is explicit.
      TOMBSTONE -- mark slot ``i`` invalid (v=0); content stays until EVICT.
      EVICT     -- free slot ``i`` entirely.
      NOOP      -- nothing.
    """
    new_ledger = [Slot(**vars(s)) for s in ledger]
    stats = {"ADD": 0, "UPDATE": 0, "TOMBSTONE": 0, "EVICT": 0,
             "NOOP": 0, "dropped": 0}
    for op in ops:
        kind = op["op"]
        stats[kind] += 1
        if kind == "NOOP":
            continue
        if kind == "ADD":
            target = next((s for s in new_ledger if not s.occupied), None)
            if target is None:
                stats["dropped"] += 1
                continue
            target.occupied = True
            target.valid = True
            target.t_last = session_idx
            target.n_uses = 0
            target.subject = str(op.get("s", ""))
            target.predicate = str(op.get("p", ""))
            target.object = str(op.get("o", ""))
        elif kind == "UPDATE":
            i = int(op.get("i", -1))
            if not (0 <= i < len(new_ledger)) or not new_ledger[i].occupied:
                stats["dropped"] += 1
                continue
            s = new_ledger[i]
            if op.get("s") is not None:
                s.subject = str(op["s"])
            if op.get("p") is not None:
                s.predicate = str(op["p"])
            s.object = str(op.get("o", s.object))
            s.t_last = session_idx
            s.valid = True
        elif kind == "TOMBSTONE":
            i = int(op.get("i", -1))
            if not (0 <= i < len(new_ledger)) or not new_ledger[i].occupied:
                stats["dropped"] += 1
                continue
            new_ledger[i].valid = False
        elif kind == "EVICT":
            i = int(op.get("i", -1))
            if not (0 <= i < len(new_ledger)) or not new_ledger[i].occupied:
                stats["dropped"] += 1
                continue
            new_ledger[i] = Slot(id=i)
    return new_ledger, stats


def salience(slot: Slot, now: int) -> float:
    """Gold salience score used by the generator's eviction oracle.

    Frequently-queried, currently-valid, recently-written facts are kept;
    old, never-queried, tombstoned facts are evicted first. This is a
    documented oracle, not a model component: the generator needs a
    deterministic ground truth for eviction supervision.
    """
    age = max(0, now - slot.t_last)
    return 4.0 * slot.n_uses + (3.0 if slot.valid else 0.0) + max(0.0, 4.0 - age)


def enforce_budget(
    ledger: List[Slot], ops: List[Dict[str, Any]], now: int
) -> Tuple[List[Slot], List[Dict[str, Any]]]:
    """Apply ops, then evict lowest-salience occupied slots until the
    occupied count fits the budget (== number of slots). Appends EVICT ops
    to the program so the policy learns budget-aware behavior."""
    ledger, _ = apply_ops(ledger, ops, now)
    ops = list(ops)
    occupied = [s for s in ledger if s.occupied]
    for _ in range(max(0, len(occupied) - len(ledger))):
        victim = min((s for s in ledger if s.occupied),
                     key=lambda s: salience(s, now))
        ops.append({"op": "EVICT", "i": victim.id})
        ledger[victim.id] = Slot(id=victim.id)
    return ledger, ops


def facts_of(ledger: List[Slot], valid_only: bool = True) -> Dict[Tuple[str, str], str]:
    """Current fact map {(subject, predicate): object}.

    valid_only=True keeps only active slots; False also returns tombstoned
    slots (last-known value), which backs temporal questions."""
    out: Dict[Tuple[str, str], str] = {}
    for s in ledger:
        if s.occupied and (s.valid or not valid_only):
            out[(s.subject, s.predicate)] = s.object
    return out


def ledger_stats(ledger: List[Slot]) -> Dict[str, int]:
    occ = [s for s in ledger if s.occupied]
    return {
        "occupied": len(occ),
        "valid": sum(1 for s in occ if s.valid),
        "tombstoned": sum(1 for s in occ if not s.valid),
        "budget": len(ledger),
    }