"""
CASE-STUDY-03 ingest: Mooncake production traces -> engine measurements.

Provenance: github.com/kvcache-ai/Mooncake, FAST25-release/traces/ (Qin et al.,
arXiv 2407.00079 / FAST'25; Apache-2.0). Pinned SHA-256 — refuses non-matching bytes:
  conversation_trace.jsonl  b8cbb061a85206d729d91cdc2981f43c9e0d99209dce588d3af5f7934408b9df
  toolagent_trace.jsonl     48a2db1a13d3bc05e6330140c64f604ba366df20d3c9e128b5c35a01c1fa5f71
NEVER use arxiv-trace/mooncake_trace.jsonl: same row count as toolagent (23,608),
different bytes — an earlier release of the same workload.

Schema per line: {"timestamp": ms, "input_length": tokens, "output_length": tokens,
"hash_ids": [ints]} — hash_ids are remapped prefix-block hashes, BLOCK_TOKENS=512
per block with the final block of each request possibly partial.

B59: load() also accepts the portable lgw-trace/1 format (hasher/TRACE-FORMAT.md):
line 1 metadata declares hash_granularity and timestamp_unit; rows carry
prefix_hashes, normalised on load to the internal hash_ids shape. Mooncake files
declare nothing, so their granularity (512) is supplied explicitly by ingest().
Granularity has NO default anywhere in this module.

TWO REUSE DEFINITIONS (both computed; formulas in the report):
  A block-occurrence share: reused block occurrences / all block occurrences.
  B token-weighted share:  input tokens lying in previously-seen blocks / all input
    tokens, where block i of a request carries min(512, input_length - 512*i) tokens.
Definition B is what a prefix cache saves money on and is what feeds the engine as
f_prefix. A sits BELOW B whenever reuse concentrates in early, full-size blocks.

h_target comes from an LRU cache-policy REPLAY at stated capacities AND A STATED TTL
(block-granular LRU; touch-on-hit, insert-on-miss, lazy expiry), token-weighted hit
rate. The full capacity curve is published; the engine input takes the stated
REFERENCE_CAPACITY point at DEFAULT_TIER. AUDIT-06 C1: an untimed replay on a
59-minute trace IS a 1-hour-TTL replay, and pricing it at the 5-minute write multiple
describes a configuration no card sells — hence TTL_TIERS, which binds each retention
window to the write multiple that tier is priced at.

RULE 9: verification scope derived from THIS artifact's input space —
test_mooncake_ingest.py (fixtures incl. a hand-derivable LRU case). Engine, library,
calculator untouched (rule 10).

Run: python3 mooncake_ingest.py <dir-with-jsonl> [out.json]
"""
from __future__ import annotations

import collections
import hashlib
import json
import pathlib
import sys

FILES = {
    "conversation_trace.jsonl": "b8cbb061a85206d729d91cdc2981f43c9e0d99209dce588d3af5f7934408b9df",
    "toolagent_trace.jsonl": "48a2db1a13d3bc05e6330140c64f604ba366df20d3c9e128b5c35a01c1fa5f71",
}
BLOCK_TOKENS = 512
CAPACITIES = [2_000, 8_000, 32_000, 128_000, None]   # blocks; None = unbounded ceiling
REFERENCE_CAPACITY = 32_000                          # stated engine input point (~16M tokens)

# AUDIT-06 C1: a replay without a TTL is a 1-hour-TTL replay on a 59-minute trace.
# Each tier pairs the retention window it replays with the cache-WRITE multiple that
# tier is actually priced at on the cited card — the pairing is the whole point:
# 1-hour retention at 5-minute write prices is not a purchasable configuration.
TTL_TIERS = {
    "5min":  {"ttl_seconds": 300,  "w": 1.25, "label": "5-minute cache (default tier)"},
    "1hour": {"ttl_seconds": 3600, "w": 2.00, "label": "1-hour cache (extended tier)"},
}
# Retention WINDOWS, measured independently of any vendor's price card. TTL_TIERS
# binds Anthropic's two PURCHASABLE pairings; this probe measures the achievable hit
# rate at the windows other published cards document, so a card is never priced
# against a hit rate measured at somebody else's window. That is the generalisation of
# the AUDIT-06 C1 defect: v1.0.0 paired 1-hour retention with 5-minute write pricing.
# Deliberately a bare list of seconds — attaching a write multiple here would recreate
# exactly the coupling that produced that error.
#   300  Anthropic default; also the short end of OpenAI's pre-5.6 idle window
#   600  the long end of OpenAI's pre-5.6 "5 to 10 minutes of inactivity"
#   1800 OpenAI GPT-5.6+ "at least 30 minutes"
#   3600 Anthropic extended tier; also the pre-5.6 stated maximum
RETENTION_PROBE_SECONDS = (300, 600, 1800, 3600)
DEFAULT_TIER = "5min"          # the stated operating point for the engine input
REFRESH_ON_HIT = True          # a cache read refreshes the entry; hard expiry = the
                               # pessimistic edge of the published sensitivity band
# Pricing: Anthropic rate card as of 2026-07-29 (platform.claude.com prompt-caching
# docs): Claude Sonnet 5 $2/1M input, $10/1M output — announced as introductory
# pricing through Aug 31, 2026 and confirmed the standard rate on 2026-08-25
# (ops/CARD-READ-2026-08-25.md); the scheduled $3/$15 increase will not occur;
# cache write 1.25x, cache read 0.1x — the library's r/w vendor constants.
P_IN, P_OUT = 2.0, 10.0
UNIT_REQUESTS = 1_000_000   # unit-economics profile: spend == $ per 1M requests


def sha256(path: pathlib.Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def block_tokens(input_length: int, n_blocks: int, i: int, *, bt: int) -> int:
    """Length units carried by block i of a request (final block may be partial).

    bt is the trace's declared granularity (B59): keyword-required, NO default —
    an attribution at the wrong block size produces a plausible wrong number.
    """
    if i < n_blocks - 1:
        return bt
    return max(0, input_length - bt * (n_blocks - 1))


_LGW_FORMAT = "lgw-trace/1"
_MOONCAKE_KEYS = {"timestamp", "input_length", "output_length", "hash_ids"}
_BOTH_SCHEMAS = (
    f"accepted schemas: {_LGW_FORMAT} (line 1 = a metadata object declaring "
    f'"format": "{_LGW_FORMAT}", hash_granularity and timestamp_unit — see '
    "hasher/TRACE-FORMAT.md) or mooncake (every line = {timestamp, input_length, "
    "output_length, hash_ids}, no metadata — granularity must be supplied by the "
    "caller)")


def _validate_row(r, ln, bt, ids_key, ids_type, prev_ts):
    """One row of either schema, validated with the trace's own granularity bt."""
    if set(r) != {"timestamp", "input_length", "output_length", ids_key}:
        raise SystemExit(f"SCHEMA MISMATCH line {ln}: keys {sorted(r)}")
    if not (isinstance(r["input_length"], int) and isinstance(r["output_length"], int)
            and r["input_length"] >= 0 and r["output_length"] >= 0):
        raise SystemExit(f"BAD TOKEN COUNTS line {ln}")
    ids = r[ids_key]
    if not (isinstance(ids, list) and all(isinstance(h, ids_type) for h in ids)):
        raise SystemExit(f"BAD {ids_key} line {ln}")
    if ids and r["input_length"] > bt * len(ids):
        raise SystemExit(f"BLOCK ACCOUNTING VIOLATION line {ln}: "
                         f"{r['input_length']} > {bt}*{len(ids)}")
    if prev_ts is not None and r["timestamp"] < prev_ts:
        raise SystemExit(f"NON-MONOTONIC timestamp line {ln}")


def load(path: pathlib.Path, *, block_tokens: int | None = None) -> tuple[list[dict], dict]:
    """Load a request trace in either accepted schema; return (rows, meta).

    Dispatch is on a DECLARED discriminator, never a try/except fallback (B59): a
    first line carrying "format" is lgw-trace/1 metadata; a first line carrying
    exactly the mooncake keys is a mooncake data row; anything else is refused
    naming both schemas. Granularity has NO default: lgw-trace/1 declares
    hash_granularity in its metadata (block_tokens, if also supplied, must agree);
    mooncake files declare nothing, so block_tokens is required (the CS-03
    reproduction supplies 512). Declared units are validated, not inferred or
    converted — both accepted schemas are millisecond-timestamped.
    rows come back in the internal shape {timestamp, input_length, output_length,
    hash_ids}; meta carries schema, block_tokens, timestamp_unit, lengths_unit.
    """
    fh = open(path, encoding="utf-8")
    first_line = fh.readline()
    try:
        first = json.loads(first_line)
    except ValueError:
        raise SystemExit(f"UNRECOGNISED TRACE {path}: line 1 is not JSON; {_BOTH_SCHEMAS}")
    rows: list[dict] = []
    prev_ts = None
    if isinstance(first, dict) and "format" in first:
        if first["format"] != _LGW_FORMAT:
            raise SystemExit(f"UNSUPPORTED FORMAT {first['format']!r} line 1; {_BOTH_SCHEMAS}")
        gran = first.get("hash_granularity")
        if not (isinstance(gran, int) and gran >= 1):
            raise SystemExit(f"MISSING GRANULARITY: {_LGW_FORMAT} metadata must declare "
                             f"hash_granularity as an int >= 1; got {gran!r}")
        if block_tokens is not None and block_tokens != gran:
            raise SystemExit(f"GRANULARITY CONFLICT: caller supplied block_tokens="
                             f"{block_tokens} but the trace metadata declares "
                             f"hash_granularity={gran}; drop the argument or fix the trace")
        unit = first.get("timestamp_unit")
        if unit != "ms":
            raise SystemExit(f"UNSUPPORTED timestamp_unit {unit!r}: declared units are "
                             "validated, not converted; supported: 'ms'")
        for ln, line in enumerate(fh, 2):
            r = json.loads(line)
            _validate_row(r, ln, gran, "prefix_hashes", str, prev_ts)
            prev_ts = r["timestamp"]
            rows.append({"timestamp": r["timestamp"], "input_length": r["input_length"],
                         "output_length": r["output_length"], "hash_ids": r["prefix_hashes"]})
        return rows, {"schema": _LGW_FORMAT, "block_tokens": gran, "timestamp_unit": "ms",
                      "lengths_unit": first.get("lengths_unit", "chars")}
    if isinstance(first, dict) and set(first) == _MOONCAKE_KEYS:
        if block_tokens is None:
            raise SystemExit("MISSING GRANULARITY: a mooncake-schema trace carries no "
                             "metadata, so load() requires block_tokens=<length units "
                             "per block> explicitly (the CS-03 reproduction supplies "
                             "512); there is no default")
        _validate_row(first, 1, block_tokens, "hash_ids", int, None)
        prev_ts = first["timestamp"]
        rows.append(first)
        for ln, line in enumerate(fh, 2):
            r = json.loads(line)
            _validate_row(r, ln, block_tokens, "hash_ids", int, prev_ts)
            prev_ts = r["timestamp"]
            rows.append(r)
        return rows, {"schema": "mooncake", "block_tokens": block_tokens,
                      "timestamp_unit": "ms", "lengths_unit": "tokens"}
    raise SystemExit(f"UNRECOGNISED TRACE {path}: line 1 matches neither schema (got "
                     f"{sorted(first) if isinstance(first, dict) else type(first).__name__}); "
                     f"{_BOTH_SCHEMAS}")


def reuse_both_definitions(rows: list[dict], *, bt: int) -> dict:
    seen: set[int | str] = set()
    occ_total = occ_reused = 0
    tok_total = tok_reused = 0
    for r in rows:
        n = len(r["hash_ids"])
        for i, h in enumerate(r["hash_ids"]):
            t = block_tokens(r["input_length"], n, i, bt=bt)
            occ_total += 1
            tok_total += t
            if h in seen:
                occ_reused += 1
                tok_reused += t
            else:
                seen.add(h)
    return {"A_block_occurrence_share": occ_reused / occ_total,
            "B_token_weighted_share": tok_reused / tok_total,
            "unique_blocks": len(seen), "block_occurrences": occ_total,
            "input_tokens": tok_total}


def lru_replay(rows: list[dict], capacity_blocks: int | None,
               ttl_seconds: float | None = None, refresh_on_hit: bool = REFRESH_ON_HIT,
               contiguous: bool = False, *, bt: int) -> dict:
    """Block-granular LRU with optional TTL expiry.

    ttl_seconds=None reproduces the pre-AUDIT-06 capacity-only behaviour (retained so
    the unbounded run still equals the definition-A/B ceiling exactly — rule 10).
    refresh_on_hit=True models a cache read extending the entry's life (the vendor's
    "hits & refreshes" wording); False is hard expiry from write time, the pessimistic
    edge of the sensitivity band.
    contiguous=True serves only the unbroken prefix from position 0 — what a real
    prefix cache can reuse (definition C; see report section on reuse definitions).
    bt is the trace's declared block granularity (B59): keyword-required, NO default.
    """
    cache: collections.OrderedDict[int | str, float | None] = collections.OrderedDict()
    occ = occ_hit = tok = tok_hit = 0
    for r in rows:
        now = r["timestamp"]
        n = len(r["hash_ids"])
        broken = False
        for i, h in enumerate(r["hash_ids"]):
            t = block_tokens(r["input_length"], n, i, bt=bt)
            occ += 1
            tok += t
            live = h in cache and (ttl_seconds is None or cache[h] > now)
            if h in cache and ttl_seconds is not None and cache[h] <= now:
                del cache[h]                       # lazily expire
            if contiguous and broken:
                live = False                       # prefix chain already broken
            if live:
                occ_hit += 1
                tok_hit += t
                cache.move_to_end(h)
                if ttl_seconds is not None and refresh_on_hit:
                    cache[h] = now + ttl_seconds * 1000
            else:
                broken = True
                cache[h] = None if ttl_seconds is None else now + ttl_seconds * 1000
                cache.move_to_end(h)
                if capacity_blocks is not None and len(cache) > capacity_blocks:
                    cache.popitem(last=False)
    return {"capacity_blocks": capacity_blocks,
            "capacity_tokens": None if capacity_blocks is None else capacity_blocks * bt,
            "ttl_seconds": ttl_seconds, "refresh_on_hit": refresh_on_hit,
            "contiguous": contiguous,
            "hit_rate_tokens": tok_hit / tok, "hit_rate_occurrences": occ_hit / occ}


def reuse_distance(rows: list[dict]) -> dict:
    """AUDIT-06 item 8: how far apart reuses happen — the statistic that makes the
    TTL result legible (a 5-minute window cannot serve a 40-minute-old prefix)."""
    last: dict[int | str, float] = {}
    gaps: list[float] = []
    for r in rows:
        for h in r["hash_ids"]:
            if h in last:
                gaps.append((r["timestamp"] - last[h]) / 1000.0)
            last[h] = r["timestamp"]
    gaps.sort()
    n = len(gaps)
    return {"reuse_events": n, "median_gap_s": gaps[n // 2], "p90_gap_s": gaps[int(n * 0.9)],
            "max_gap_s": gaps[-1],
            "share_beyond_5min": sum(1 for g in gaps if g > 300) / n}


def duplicates(rows: list[dict]) -> dict:
    """AG-03: repeats of an already-seen (hash_ids, input_length) request."""
    seen: set[tuple] = set()
    dup_calls = 0
    dup_tokens = 0
    all_tokens = 0
    dup_in = 0
    all_in = 0
    for r in rows:
        key = (tuple(r["hash_ids"]), r["input_length"])
        t = r["input_length"] + r["output_length"]
        all_tokens += t
        all_in += r["input_length"]
        if key in seen:
            dup_calls += 1
            dup_tokens += t
            dup_in += r["input_length"]
        else:
            seen.add(key)
    # input_token_share is the magnitude that belongs in the CA-01 overlap disclosure:
    # f_prefix is an input-side measure, so the overlapping quantity is input tokens.
    return {"call_share": dup_calls / len(rows), "token_share": dup_tokens / all_tokens,
            "input_token_share": dup_in / all_in, "duplicate_calls": dup_calls}


def method_demos(rows: list[dict], *, bt: int) -> dict:
    """RS-10 / CA-04 method demonstrations (config-required vars stay unmeasured)."""
    freq: collections.Counter = collections.Counter()
    for r in rows:
        for h in set(r["hash_ids"]):
            freq[h] += 1
    n = len(rows)
    common = {h for h, c in freq.items() if c >= 0.5 * n}
    tok_common = tok_total = 0
    for r in rows:
        nb = len(r["hash_ids"])
        for i, h in enumerate(r["hash_ids"]):
            t = block_tokens(r["input_length"], nb, i, bt=bt)
            tok_total += t
            if h in common:
                tok_common += t
    return {"rs10_universal_block_token_share": tok_common / tok_total if tok_total else 0.0,
            "rs10_universal_blocks": len(common)}


def ingest(data_dir: pathlib.Path) -> dict:
    out = {"provenance": {
               "source": "github.com/kvcache-ai/Mooncake FAST25-release/traces "
                         "(Qin et al., arXiv 2407.00079; Apache-2.0)",
               "excluded": "arxiv-trace/mooncake_trace.jsonl (earlier release, same "
                           "row count as toolagent, different bytes — never combine)",
               "files": {}},
           "pricing": {"card": "Anthropic Claude Sonnet 5, platform.claude.com, "
                               "as of 2026-07-29, re-read 2026-08-25: now the "
                               "standard rate, the scheduled increase will not occur",
                       "p_in_per_1m": P_IN, "p_out_per_1m": P_OUT,
                       "cache_read_multiple_r": 0.1, "cache_write_multiple_w": 1.25},
           "traces": {}}
    for fname, want in FILES.items():
        path = data_dir / fname
        got = sha256(path)
        if got != want:
            raise SystemExit(f"HASH MISMATCH {fname}: {got}")
        rows, _tmeta = load(path, block_tokens=BLOCK_TOKENS)
        span_s = (rows[-1]["timestamp"] - rows[0]["timestamp"]) / 1000.0
        tin = sum(r["input_length"] for r in rows)
        tout = sum(r["output_length"] for r in rows)
        reuse = reuse_both_definitions(rows, bt=BLOCK_TOKENS)
        # AUDIT-06 V3 / definition C: contiguity-only reuse — what a real prefix cache
        # can serve. Published because it is the first objection to block-independent
        # counting, and on this data C == B exactly (zero hits occur after a miss).
        hits_after_miss = 0
        total_hits = 0
        seen_c: set[int] = set()
        for rr in rows:
            broke = False
            for h in rr["hash_ids"]:
                if h in seen_c:
                    total_hits += 1
                    if broke:
                        hits_after_miss += 1
                else:
                    seen_c.add(h)
                    broke = True
        contig = lru_replay(rows, None, None, contiguous=True, bt=BLOCK_TOKENS)
        reuse["C_contiguous_token_share"] = contig["hit_rate_tokens"]
        reuse["hits_after_miss"] = hits_after_miss
        reuse["total_hits"] = total_hits
        reuse["C_equals_B"] = abs(contig["hit_rate_tokens"] - reuse["B_token_weighted_share"]) < 1e-12

        default = TTL_TIERS[DEFAULT_TIER]
        # capacity curve at the DEFAULT operating tier (what the engine input uses)
        curve = [lru_replay(rows, c, default["ttl_seconds"], bt=BLOCK_TOKENS) for c in CAPACITIES]
        # untimed ceiling curve retained for reference (equals the definition-A/B
        # ceiling when unbounded — rule 10 property, asserted below)
        curve_no_ttl = [lru_replay(rows, c, bt=BLOCK_TOKENS) for c in CAPACITIES]
        ref = next(p for p in curve if p["capacity_blocks"] == REFERENCE_CAPACITY)
        tiers = {}
        for tname, t in TTL_TIERS.items():
            tiers[tname] = {
                "ttl_seconds": t["ttl_seconds"], "w_write_multiple": t["w"],
                "label": t["label"],
                "h_target_refresh": lru_replay(rows, REFERENCE_CAPACITY, t["ttl_seconds"],
                                               True, bt=BLOCK_TOKENS)["hit_rate_tokens"],
                "h_target_hard_expiry": lru_replay(rows, REFERENCE_CAPACITY, t["ttl_seconds"],
                                                   False, bt=BLOCK_TOKENS)["hit_rate_tokens"],
                "break_even_h": (t["w"] - 1) / (t["w"] - 0.1),
            }
        probe = {str(sec): lru_replay(rows, REFERENCE_CAPACITY, sec,
                                      REFRESH_ON_HIT, bt=BLOCK_TOKENS)["hit_rate_tokens"]
                 for sec in RETENTION_PROBE_SECONDS}
        # rule 10, asserted on every run: the probe and the tier table must be the same
        # replay, not two implementations that happen to agree today.
        for tname, t in TTL_TIERS.items():
            k = str(t["ttl_seconds"])
            if k in probe:
                assert abs(probe[k] - tiers[tname]["h_target_refresh"]) < 1e-12, \
                    f"retention probe disagrees with tier {tname}"
        assert all(probe[str(b)] >= probe[str(a)] - 1e-12
                   for a, b in zip(RETENTION_PROBE_SECONDS, RETENTION_PROBE_SECONDS[1:])), \
            "retention probe is not monotone in window length"
        assert abs(curve_no_ttl[-1]["hit_rate_occurrences"]
                   - reuse["A_block_occurrence_share"]) < 1e-12, "rule 10: ceiling drift"
        dup = duplicates(rows)
        out["provenance"]["files"][fname] = {"sha256": want, "rows": len(rows)}
        out["traces"][fname.replace("_trace.jsonl", "")] = {
            "rows": len(rows), "span_seconds": span_s,
            "mean_in": tin / len(rows), "mean_out": tout / len(rows),
            "reuse": reuse, "lru_curve": curve, "lru_curve_untimed": curve_no_ttl,
            "reuse_distance": reuse_distance(rows),
            "tiers": tiers, "default_tier": DEFAULT_TIER,
            "retention_probe": probe,
            "reference_capacity_blocks": REFERENCE_CAPACITY,
            "duplicates": dup, "method_demos": method_demos(rows, bt=BLOCK_TOKENS),
            "profile": {"N": float(UNIT_REQUESTS), "L": 1.0,
                        "T_in": tin / len(rows), "T_out": tout / len(rows),
                        "p_in": P_IN, "p_out": P_OUT,
                        "S_aux": 0.0, "S_infra": 0.0, "stack": "api_only"},
            "overrides": {"f_prefix": reuse["B_token_weighted_share"],
                          "CA-01.h_target": ref["hit_rate_tokens"],  # DEFAULT_TIER
                          "AG-03.pct_duplicate_calls": dup["token_share"]},
        }
    return out


def main() -> None:
    data_dir = pathlib.Path(sys.argv[1])
    dest = (pathlib.Path(sys.argv[2]) if len(sys.argv) > 2
            else pathlib.Path(__file__).parent / "measurements_cs03.json")
    out = ingest(data_dir)
    with dest.open("w", encoding="utf-8", newline="\n") as _f:
        # B59 writer alignment: the committed measurements_cs03.json (the published
        # anchor; last serialised by 4f3a6c8's hand edit) carries raw UTF-8 and a
        # trailing newline. The writer moves to the anchor, never the reverse.
        _f.write(json.dumps(out, indent=1, ensure_ascii=False) + "\n")
    for name, t in out["traces"].items():
        r = t["reuse"]
        print(f"{name}: rows={t['rows']:,} A={r['A_block_occurrence_share']:.3f} "
              f"B={r['B_token_weighted_share']:.3f} C==B={r['C_equals_B']} "
              f"h(5min)={t['tiers']['5min']['h_target_refresh']:.4f} "
              f"h(1hour)={t['tiers']['1hour']['h_target_refresh']:.4f} "
              f"dup_tok={t['duplicates']['token_share']:.4f}")


if __name__ == "__main__":
    main()