"""Full installed MTPLX forward oracle; run only under test-supervisor.

Uses the exact DS4Server native canary entry point from a linked bridge dylib.
Diagnostic, not a clean generation benchmark: optional original AR to EOS.
"""
import ctypes
import argparse
import hashlib
import json
import os
from pathlib import Path
import sys
import threading
import time


def ar_oracle(model, tokenizer, ids, root, set_phase):
    """Actual _prefill and generate_ar, natural EOS, no tools or downloads."""
    import mlx.core as mx
    import numpy as np
    from mtplx.runtime import MTPLXRuntime
    from mtplx.sampling import SamplerConfig
    from mtplx.generation import (_prefill, generate_ar, _final_logits_prefill_enabled,
        _sustained_prefill_enabled, _prefill_chunk_size)
    rt = MTPLXRuntime(model, tokenizer, root, False, None)
    def fingerprint(a):
        values = np.asarray(a.astype(mx.float32)).astype('<f4')
        assert np.isfinite(values).all()
        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
    set_phase("prefill")
    cache, logits, hidden, elapsed = _prefill(rt, ids, return_hidden=False)
    print(json.dumps(dict(event="ar_prefill_reference", token_ids=ids, logits=fingerprint(logits),
        hidden=hidden is not None, seconds=elapsed, chunks=rt.diagnostic_counters.get("prefill_chunks", 0),
        chunk_size=_prefill_chunk_size() if _sustained_prefill_enabled() else None,
        final_logits_only=_final_logits_prefill_enabled())), flush=True)
    del cache, logits, hidden
    sampler = SamplerConfig(temperature=.6, top_p=.95, top_k=20)
    def progress(data):
        set_phase("decode" if data["phase"] == "completed" else "prefill")
        print(json.dumps(dict(event="ar_prefill_progress", **data)), file=sys.stderr, flush=True)
    out = generate_ar(rt, ids, max_tokens=2**63-1, sampler=sampler, seed=0,
        stop_token_ids=set(tokenizer.eos_token_ids), capture_final_state=True,
        prefill_callback=progress,
        token_callback=lambda ids: print(json.dumps(dict(event="ar_token", ids=ids)), file=sys.stderr, flush=True))
    print(json.dumps(dict(event="ar_reference_end", tokens=out.tokens, text=out.text,
        finish_reason=out.finish_reason, stops=sorted(tokenizer.eos_token_ids), seed=0,
        temperature=.6, top_p=.95, top_k=20, verify_calls=out.stats.verify_calls,
        elapsed_seconds=out.stats.elapsed_s, prompt_seconds=out.stats.prompt_eval_time_s,
        final_logits=None if out.final_state is None else fingerprint(out.final_state.final_logits))), flush=True)


def mtp_prefill_oracle(model, tokenizer, ids, root, set_phase):
    """Original streaming history prefill and the first actual draft, no generation claim."""
    import mlx.core as mx
    import numpy as np
    from mtplx.runtime import MTPLXRuntime
    from mtplx.mtp_patch import MTPContract
    from mtplx.generation import _prefill_committed_mtp_history_streaming, _prefill_chunk_size, _final_logits_prefill_enabled
    rt = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
    def fingerprint(a):
        values = np.asarray(a.astype(mx.float32)).astype('<f4')
        assert np.isfinite(values).all()
        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
    for window in (None, 3):
        set_phase("prefill")
        cache, logits, hidden, history, target_s, history_s, base = _prefill_committed_mtp_history_streaming(
            rt, ids, history_window_tokens=window,
            chunk_callback=lambda data: print(json.dumps(dict(event="mtp_prefill_progress", **data)), file=sys.stderr, flush=True))
        receipt = dict(event="mtp_prefill_reference", token_ids=ids, window=window,
            chunk_size=_prefill_chunk_size(), final_only=_final_logits_prefill_enabled(),
            logits=fingerprint(logits), hidden=fingerprint(hidden), offset=history[0].offset,
            base=base, target_seconds=target_s, history_seconds=history_s)
        set_phase("decode")
        token = int(mx.argmax(logits[0]).item())
        draft_logits, draft_hidden = model.mtp_forward(hidden, mx.array([[token]]), mtp_cache=history, return_hidden=True)
        mx.eval(draft_logits, draft_hidden)
        receipt.update(token=token, draft_logits=fingerprint(draft_logits), draft_hidden=fingerprint(draft_hidden), draft_offset=history[0].offset)
        print(json.dumps(receipt), flush=True)
        del cache, logits, hidden, history, draft_logits, draft_hidden
    print(json.dumps(dict(event="mtp_prefill_reference_end")), flush=True)


def mtp_generation_oracle(model, tokenizer, ids, root, set_phase, restored_suffixes=None, snapshot_restore=False):
    """Actual generate_mtpk to natural EOS; preserves standard family p/q policy."""
    import mlx.core as mx
    import numpy as np
    from mtplx.runtime import MTPLXRuntime
    from mtplx.mtp_patch import MTPContract
    from mtplx.sampling import SamplerConfig
    from mtplx.generation import generate_mtpk, _prefill_chunk_size, _final_logits_prefill_enabled
    rt = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
    def fingerprint(a):
        values = np.asarray(a.astype(mx.float32)).astype('<f4')
        assert np.isfinite(values).all()
        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
    sampler = SamplerConfig(temperature=.6, top_p=.95, top_k=20)
    def progress(data):
        set_phase("decode" if data["phase"] == "completed" else "prefill")
        print(json.dumps(dict(event="mtp_prefill_progress", **data)), file=sys.stderr, flush=True)
    print(json.dumps(dict(event="mtp_generation_input", token_ids=ids, temperature=.6, top_p=.95, top_k=20,
        seed=0, depth=3, stops=sorted(tokenizer.eos_token_ids), history="committed",
        restored_suffix_cases=3 if restored_suffixes is not None else 0,
        chunk_size=_prefill_chunk_size(), final_only=_final_logits_prefill_enabled())), flush=True)
    out = generate_mtpk(rt, ids, max_tokens=2**63-1, sampler=sampler, speculative_depth=3,
        seed=0, stop_token_ids=set(tokenizer.eos_token_ids), mtp_history_policy="committed",
        verify_strategy="batched", capture_final_state=True, prefill_callback=progress,
        token_callback=lambda ids: print(json.dumps(dict(event="mtp_tokens", ids=ids)), file=sys.stderr, flush=True))
    cycles = [dict(primary=e["primary"], primary_already_emitted=e["primary_already_emitted"],
        drafts=[d["token"] for d in e["drafts"]], accepted=e["accepted_depths"],
        correction=next((d["correction"] for d in e["drafts"] if d.get("accepted") is False), None),
        bonus=e.get("bonus_token"), capture_committed=str(e.get("capture_repair", "")).startswith("captured_prefix"),
        **({"context_copy": {k: v for k, v in e["context_copy"].items() if k != "time_s"}}
           if "context_copy" in e else {}))
        for e in out.stats.events if "primary" in e]
    final = out.final_state
    assert final is not None
    print(json.dumps(dict(event="mtp_generation_end", tokens=out.tokens, text=out.text, cycles=cycles,
        raw_cycles=out.stats.events,
        finish_reason=out.finish_reason, safe_to_commit=final.safe_to_commit,
        logits=fingerprint(final.final_logits), hidden=fingerprint(final.final_hidden),
        history_offset=final.final_committed_mtp_cache[0].offset,
        verify_calls=out.stats.verify_calls, accepted=out.stats.accepted_drafts, drafted=out.stats.drafted_tokens,
        elapsed_seconds=out.stats.elapsed_s, prefill_seconds=out.stats.prompt_eval_time_s)), flush=True)
    if restored_suffixes is not None:
        from types import SimpleNamespace
        from mtplx.generation import _prefill_restored_prompt_suffix, _small_suffix_fused_max, _runtime_counter_snapshot, _runtime_counter_delta
        from mtplx.cache_state import snapshot_cache, snapshot_cache_lazy_hybrid, restore_cache
        def snapshot_digest(snapshot):
            h = hashlib.sha256()
            for state in snapshot.states:
                h.update(b"[")
                for a in state:
                    if a is None:
                        h.update(b"N")
                    else:
                        tag = {mx.bfloat16:"BF16", mx.float32:"F32", mx.int32:"I32", mx.int64:"I64", mx.uint32:"U32", mx.float16:"F16"}[a.dtype]
                        h.update(tag.encode())
                        h.update(json.dumps(list(a.shape), separators=(",", ":")).encode())
                        h.update(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes())
                h.update(b"]")
            return h.hexdigest()
        restored = SimpleNamespace(cache=final.final_trunk_cache,
            mtp_history_cache=final.final_committed_mtp_cache, hidden=final.final_hidden)
        cached = len(ids) + len(out.tokens)
        # Numerical warm-extension cases on the same live final cache, not a
        # substitute for the full session-bank/ongoing-chat acceptance run.
        texts = ["Tell me a complete short story about a lighthouse keeper. Do not ask questions.",
            "Give a summary of the following text:\n\n" + Path(restored_suffixes).read_text(),
            "Write a Python function is_prime(n: int) -> bool, followed by five assert examples. No tools."]
        for step, text in enumerate(texts):
            suffix = tokenizer.encode("\n<|im_start|>user\n" + text + "<|im_end|>\n<|im_start|>assistant\n<think>\n")
            chunks = []
            set_phase("prefill")
            saved = None
            if snapshot_restore:
                lazy_kv = step != 1
                capture = snapshot_cache_lazy_hybrid if lazy_kv else snapshot_cache
                trunk_snapshot, head_snapshot = capture(restored.cache), capture(restored.mtp_history_cache)
                saved = dict(lazy_kv=lazy_kv, trunk=snapshot_digest(trunk_snapshot), mtp=snapshot_digest(head_snapshot))
                restored.cache, restored.mtp_history_cache = model.make_cache(), model.make_mtp_cache()
                restore_cache(restored.cache, trunk_snapshot, clone_states=not lazy_kv)
                restore_cache(restored.mtp_history_cache, head_snapshot, clone_states=not lazy_kv)
                print(json.dumps(dict(event="snapshot_restored", step=step, **saved)), file=sys.stderr, flush=True)
            def chunk(data):
                chunks.append({k: data[k] for k in ("tokens_done", "tokens_total", "cached_tokens", "new_prefill_tokens", "chunk_size")})
                print(json.dumps(dict(event="restored_suffix_progress", **data)), file=sys.stderr, flush=True)
            counts_before = _runtime_counter_snapshot(rt)
            logits, hidden, target_s, history_s = _prefill_restored_prompt_suffix(rt, restored, suffix,
                base_hidden_variant="post_norm", mtp_hidden_variant="post_norm", mtp_history_policy="committed",
                chunk_callback=chunk, tokens_total=cached + len(suffix), cached_tokens=cached)
            if saved is not None:
                saved.update(trunk_after=snapshot_digest(trunk_snapshot), mtp_after=snapshot_digest(head_snapshot))
                assert saved["trunk"] == saved["trunk_after"] and saved["mtp"] == saved["mtp_after"], "snapshot poisoned by continuation"
            print(json.dumps(dict(event="restored_suffix_reference", step=step, token_ids=suffix, cached_tokens=cached,
                fused_max=_small_suffix_fused_max(), chunk_size=_prefill_chunk_size(), final_only=_final_logits_prefill_enabled(),
                logits=fingerprint(logits), hidden=fingerprint(hidden), history_offset=restored.mtp_history_cache[0].offset,
                chunks=chunks, snapshot=saved, runtime_counts=_runtime_counter_delta(rt, counts_before), target_seconds=target_s, history_seconds=history_s)), flush=True)
            restored.hidden = hidden
            cached += len(suffix)
        print(json.dumps(dict(event="restored_suffix_reference_end", cases=len(texts))), flush=True)


def source_chat(records, mtp, canary):
    """Accept only the explicit, complete UI-compatible comparison contract."""
    starts = [r for r in records if r.get("event") == "start"]
    results = [r for r in records if r.get("event") == "result"]
    warmups = [r for r in records if r.get("event") == "warmup_result"]
    if len(starts) != 1:
        raise ValueError("source requires exactly one start receipt")
    start = starts[0]
    settings = start["settings"]
    if not (start["model"] == "qwen3.8-flash-next" and start["plain_chat"]
            and not start["system_prompt"] and start["canary"] == canary
            and settings["reasoning"] == "low" and settings["power_percent"] == 100
            and settings["prefill_chunk"] == 2048 and settings["min_p"] == 0
            and not settings["quality"] and settings["seed"] is not None
            and settings["acceleration"]["kind"] == "mtp"
            and settings["acceleration"]["enabled"] == mtp
            and len(start["prompts"]) == 3
            and [r["turn"] for r in results] == [1, 2, 3]
            and all(r["ok"] and r["finish_reason"] == "stop" for r in results)
            and start["warmup"]["enabled"] and len(warmups) == 1 and warmups[0]["ok"]):
        raise ValueError("source must be a complete warm Power-100 Low Qwen plain chat with matching MTP/canary and supported settings")
    return start, results, warmups[0]


def token_report(event, mode, **metadata):
    """Rate-limit IO, not inference; only newly emitted tokens mean progress."""
    total, last = 0, None
    def report(tokens):
        nonlocal total, last
        if not tokens:
            return
        total += len(tokens)
        now = time.monotonic()
        if mode == "full":
            print(json.dumps(dict(event=event, ids=tokens, **metadata)), file=sys.stderr, flush=True)
        elif last is None or now - last >= 1:
            print(json.dumps(dict(event="chat_decode_progress", scope=event,
                tokens=total, **metadata)), file=sys.stderr, flush=True)
            last = now
    return report


def chat_oracle(model, tokenizer, ids, root, set_phase, mtp, source=None, token_events="full"):
    """Three complete turns through the real SessionBank and generation entrypoints."""
    import mlx.core as mx
    import numpy as np
    from mtplx.runtime import MTPLXRuntime
    from mtplx.mtp_patch import MTPContract
    from mtplx.sampling import SamplerConfig
    from mtplx.session_bank import SessionBank, _lazy_snapshot_enabled
    from mtplx.cache_state import snapshot_cache
    from mtplx.generation import (generate_ar, generate_mtpk, _prefill_chunk_size, _sustained_prefill_enabled,
        _final_logits_prefill_enabled, _small_suffix_fused_max,
        _store_on_prefill_env_enabled, _store_on_prefill_min_suffix,
        _gdn_boundary_capture_enabled, _gdn_boundary_tail_interval, _gdn_boundary_max_count,
        restore_or_prefill_prompt_state, _resolve_runtime_base_hidden_variant)
    rt = MTPLXRuntime(model, tokenizer, root, mtp, MTPContract() if mtp else None)
    bank = SessionBank()
    settings = source[0]["settings"] if source else dict(temperature=.6, top_p=.95, top_k=20, seed=0)
    sampler = SamplerConfig(temperature=settings["temperature"], top_p=settings["top_p"], top_k=settings["top_k"])
    seed = settings["seed"]
    print(json.dumps(dict(event="chat_config", mtp=mtp, turns=3, seed=seed, token_events=token_events,
        temperature=sampler.temperature, top_p=sampler.top_p, top_k=sampler.top_k,
        warmup=source[0]["warmup"] if source else None, depth=3, stops=sorted(tokenizer.eos_token_ids),
        lazy_kv=_lazy_snapshot_enabled(), chunk_size=_prefill_chunk_size(), sustained_prefill=_sustained_prefill_enabled(),
        final_only=_final_logits_prefill_enabled(), fused_max=_small_suffix_fused_max(),
        store_on_prefill=_store_on_prefill_env_enabled(), store_min_suffix=_store_on_prefill_min_suffix(),
        boundary_capture=_gdn_boundary_capture_enabled(), boundary_tail=_gdn_boundary_tail_interval(), boundary_max=_gdn_boundary_max_count())), flush=True)
    def fingerprint(a):
        if a is None:
            return None
        values = np.asarray(a.astype(mx.float32)).astype('<f4')
        assert np.isfinite(values).all()
        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
    def progress(data):
        set_phase("decode" if data["phase"] == "completed" else "prefill")
        print(json.dumps(dict(event="chat_prefill_progress", **data)), file=sys.stderr, flush=True)
    followups = [
        "Tell me a complete short story about a lighthouse keeper. Do not ask questions.",
        "Write a Python function is_prime(n: int) -> bool, followed by five assert examples. No tools.",
    ]
    generate = generate_mtpk if mtp else generate_ar
    mode = dict(speculative_depth=3, mtp_history_policy="committed", verify_strategy="batched") if mtp else {}
    if source:
        from mtplx.server.openai import ChatMessage, _encode_messages_uncached
        start, expected_turns, expected_warmup = source
        followups = start["prompts"][1:]
        # Derive the UI's system-only bootstrap from the ORIGINAL template,
        # without copying the Low instruction or hard-coding its token count.
        rendered = tokenizer.decode(ids, skip_special_tokens=False)
        prefix = rendered.split("<|im_start|>user\n", 1)[0]
        bootstrap_ids = tokenizer.encode(prefix)
        assert bootstrap_ids and list(ids[:len(bootstrap_ids)]) == bootstrap_ids
        state = restore_or_prefill_prompt_state(rt, bootstrap_ids,
            base_hidden_variant="post_norm" if mtp else None,
            mtp_history_policy="committed" if mtp else "cycle",
            capture_hidden=mtp, session_bank=bank,
            store_prefix_snapshot=False, prefill_callback=progress)
        head = snapshot_cache(state.committed_mtp_cache) if state.committed_mtp_cache is not None else None
        assert bank.put(runtime=rt, token_ids=bootstrap_ids, cache=state.trunk_cache,
            logits=state.logits, hidden=state.hidden,
            hidden_variant=_resolve_runtime_base_hidden_variant(rt, None),
            session_id="oracle-bootstrap", mtp_history_policy="committed" if mtp else "cycle",
            mtp_history_snapshot=head, snapshot_epoch=len(bootstrap_ids),
            mtp_snapshot_epoch=len(bootstrap_ids) if head is not None else None,
            gdn_boundaries=state.gdn_boundaries) is not None
        del state, head
        warm_ids = _encode_messages_uncached(tokenizer,
            [ChatMessage(role="user", content=start["warmup"]["prompt"])],
            enable_thinking=True, reasoning_effort="low", preserve_reasoning_history=True, tools=None)
        warm = generate(rt, warm_ids, max_tokens=start["warmup"]["max_generated_tokens"],
            sampler=sampler, seed=seed, stop_token_ids=set(tokenizer.eos_token_ids),
            session_bank=bank, session_id="oracle-warmup", capture_final_state=True,
            prefill_callback=progress,
            token_callback=token_report("chat_warmup_tokens", token_events), **mode)
        visible = sum(t not in tokenizer.eos_token_ids for t in warm.tokens)
        print(json.dumps(dict(event="chat_warmup", tokens=warm.tokens, text=warm.text,
            finish_reason=warm.finish_reason, bootstrap_ids=bootstrap_ids,
            cached_tokens=warm.stats.cached_tokens)), flush=True)
        assert warm.stats.cached_tokens == len(bootstrap_ids), "warmup must reuse the UI bootstrap"
        assert warm.finish_reason == expected_warmup["finish_reason"]
        assert visible == expected_warmup["completion_tokens"], "warmup token count differs"
        del warm
        bank.clear(session_id="oracle-warmup")
    for step in range(3):
        prompt_ids = list(ids)
        limit = min(settings["max_generated_tokens"], settings["context_tokens"] - len(ids)) if source else 2**63-1
        assert limit > 0, "source prompt exhausts the context"
        kwargs = dict(max_tokens=limit, sampler=sampler, seed=seed,
            stop_token_ids=set(tokenizer.eos_token_ids), session_bank=bank, session_id="oracle-chat",
            capture_final_state=True, prefill_callback=progress,
            token_callback=token_report("chat_tokens", token_events, step=step))
        out = generate_mtpk(rt, ids, speculative_depth=3, mtp_history_policy="committed",
            verify_strategy="batched", **kwargs) if mtp else generate_ar(rt, ids, **kwargs)
        final = out.final_state
        assert out.finish_reason == "stop" and final is not None and final.safe_to_commit
        cycles = [dict(primary=e["primary"], primary_already_emitted=e["primary_already_emitted"],
            drafts=[d["token"] for d in e["drafts"]], accepted=e["accepted_depths"],
            correction=next((d["correction"] for d in e["drafts"] if d.get("accepted") is False), None),
            bonus=e.get("bonus_token"), capture_committed=str(e.get("capture_repair", "")).startswith("captured_prefix"),
            **({"context_copy": {k:v for k,v in e["context_copy"].items() if k != "time_s"}} if "context_copy" in e else {}))
            for e in out.stats.events if "primary" in e]
        ids = prompt_ids + list(out.tokens)
        # Same final-state payload and head snapshot rule as openai's commit.
        head = snapshot_cache(final.final_committed_mtp_cache) if final.final_committed_mtp_cache is not None else None
        entry = bank.put(runtime=rt, token_ids=ids, cache=final.final_trunk_cache,
            logits=final.final_logits, hidden=final.final_hidden, hidden_variant="post_norm",
            session_id="oracle-chat", mtp_history_policy="committed" if mtp else "cycle",
            mtp_history_snapshot=head, snapshot_epoch=len(ids),
            mtp_snapshot_epoch=len(ids) if head is not None else None)
        assert entry is not None
        print(json.dumps(dict(event="chat_turn", step=step, prompt_ids=prompt_ids, tokens=out.tokens,
            text=out.text, finish_reason=out.finish_reason, safe_to_commit=final.safe_to_commit,
            cached_tokens=out.stats.cached_tokens, boundaries=[b[0] for b in entry.gdn_boundaries],
            logits=fingerprint(final.final_logits),
            hidden=fingerprint(final.final_hidden), cycles=cycles,
            history_offset=final.final_committed_mtp_cache[0].offset if head is not None else None,
            prefill_seconds=out.stats.prompt_eval_time_s, elapsed_seconds=out.stats.elapsed_s,
            prefill_scope="engine_target_and_mtp_history_excluding_restore_and_checkpoint",
            mtp_history_seconds=out.stats.prompt_mtp_history_time_s,
            decode_seconds=out.stats.decode_elapsed_s, decode_tps=out.stats.decode_tok_s,
            prompt_state_seconds=out.stats.prompt_state_total_time_s,
            restore_seconds=out.stats.cache_restore_time_s)), flush=True)
        if source:
            expected = expected_turns[step]
            checks = dict(
                prompt_tokens=len(prompt_ids) == expected["tokens"]["prompt"],
                cached_tokens=out.stats.cached_tokens == expected["tokens"]["cached"],
                completion_tokens=sum(t not in tokenizer.eos_token_ids for t in out.tokens) == expected["tokens"]["completion"],
                text=out.text == (expected["reasoning"] or "") + "</think>" + expected["content"])
            print(json.dumps(dict(event="chat_source_match", step=step, checks=checks)), flush=True)
            assert all(checks.values()), f"source mismatch at turn {step}: {checks}"
        if step < len(followups):
            ids += tokenizer.encode("\n<|im_start|>user\n" + followups[step] + "<|im_end|>\n<|im_start|>assistant\n<think>\n")
    print(json.dumps(dict(event="chat_end", turns=3)), flush=True)


def main(root, bridge, decode=False, server_defaults=False, mtp=False, ar=False, readme=None, mtp_prefill=False, mtp_generate=False, restored_suffixes=None, snapshot_restore=False, chat=False, source=None, canary=True, token_events="full"):
    class Sample(ctypes.Structure):
        _fields_ = [("scheduled_seconds", ctypes.c_double), ("completed_seconds", ctypes.c_double)]

    lib = ctypes.CDLL(str(Path(bridge).resolve()))
    probe = lib.ds4_gpu_canary_probe
    probe.argtypes = [ctypes.POINTER(Sample)]
    probe.restype = ctypes.c_int
    phase = "startup"
    pending = None
    stop = threading.Event()
    samples = []

    def watch():
        while not stop.wait(0.01):
            if pending is not None and time.monotonic() - pending >= 2:
                print(json.dumps(dict(event="test_canary_abort", phase=phase, reason="GPU probe exceeds 2s")), file=sys.stderr, flush=True)
                os._exit(86)

    def probes():
        nonlocal pending
        while not stop.is_set():
            at = time.monotonic()
            probe_phase = phase
            pending = at
            sample = Sample()
            ok = probe(ctypes.byref(sample))
            pending = None
            samples.append(dict(phase=probe_phase, scheduled_ms=sample.scheduled_seconds * 1000, completed_ms=sample.completed_seconds * 1000))
            if not ok or sample.completed_seconds >= 2:
                print(json.dumps(dict(event="test_canary_abort", sample=samples[-1], ok=bool(ok))), file=sys.stderr, flush=True)
                os._exit(86)
            stop.wait(max(0, 0.1 - (time.monotonic() - at)))

    watcher = threading.Thread(target=watch, daemon=True)
    prober = threading.Thread(target=probes, daemon=True)
    if canary:
        watcher.start()
        prober.start()
    try:
        os.environ["HF_HUB_OFFLINE"] = "1"
        os.environ["TRANSFORMERS_OFFLINE"] = "1"
        print("reference: importing pinned MTPLX", file=sys.stderr, flush=True)
        import mlx.core as mx
        import numpy as np
        from mlx_lm.utils import load_model
        from mtplx.models import qwen4_exp as qwen
        from mtplx.runtime import _load_tokenizer_resilient
        from mtplx.server.openai import ChatMessage, _encode_messages_uncached
        assert mx.__version__ == "0.32.2"
        root = Path(root)
        if server_defaults:
            from mtplx.server.openai import _server_runtime_env_overrides
            overrides = _server_runtime_env_overrides(
                argparse.Namespace(model=str(root), generation_mode="mtp" if mtp or mtp_prefill or mtp_generate else "ar"), None)
            os.environ.update(overrides)
            print(json.dumps(dict(event="reference_environment", overrides=overrides)), flush=True)
        if source:
            # Match policy::turn_options, as mtplx-policy-reference.py does:
            # bounded streaming prefill, request-local chunks, final logits.
            # This also governs bootstrap and restored suffixes, not just the
            # reported chunk-size number or final-logits switch.
            os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
            print(json.dumps(dict(event="reference_request_policy", sustained_prefill=True, final_logits_only=True,
                prefill_chunk=source[0]["settings"]["prefill_chunk"])), flush=True)
        tokenizer = _load_tokenizer_resilient(root, json.loads((root / "config.json").read_text()))
        prompt = "ping" if readme is None else "Give a summary of the following text:\n\n" + Path(readme).read_text()
        if source:
            assert prompt == source[0]["prompts"][0], "README prompt differs from the source run"
        print(json.dumps(dict(event="reference_prompt", text=prompt,
            sha256=hashlib.sha256(prompt.encode()).hexdigest())), flush=True)
        ids = _encode_messages_uncached(tokenizer, [ChatMessage(role="user", content=prompt)],
            enable_thinking=True, reasoning_effort="low", preserve_reasoning_history=True, tools=None)
        phase = "loading"
        print(json.dumps(dict(event="reference_loading", token_ids=ids)), file=sys.stderr, flush=True)
        started = time.monotonic()
        model, _ = load_model(root, lazy=False, strict=True, get_model_classes=lambda **_: (qwen.Model, qwen.ModelArgs))
        model.post_weight_load(root)
        if mtp or mtp_prefill or mtp_generate:
            assert model.attach_mtp(root), "installed MTP head missing"
        loaded = time.monotonic()
        if ar or mtp_prefill or mtp_generate:
            def set_phase(value):
                nonlocal phase
                phase = value
            oracle = mtp_generation_oracle if mtp_generate else mtp_prefill_oracle if mtp_prefill else ar_oracle
            if chat:
                from mtplx.generation import prefill_chunk_size_override
                with prefill_chunk_size_override(source[0]["settings"]["prefill_chunk"] if source else None):
                    chat_oracle(model, tokenizer, ids, root, set_phase, mtp_generate, source, token_events)
            elif mtp_generate:
                oracle(model, tokenizer, ids, root, set_phase, restored_suffixes, snapshot_restore)
            else:
                oracle(model, tokenizer, ids, root, set_phase)
            phase = "finishing"
            del model
            return
        print("reference: trunk loaded, starting prefill", file=sys.stderr, flush=True)
        phase = "prefill"
        cache = model.make_cache()
        prefill_ids = ids[:-1] if mtp else ids
        out = model(mx.array([prefill_ids], dtype=mx.int32), cache=cache)
        mx.eval(out)
        finished = time.monotonic()
        values = np.asarray(out.astype(mx.float32)).astype('<f4')
        assert np.isfinite(values).all()
        print(json.dumps(dict(event="forward_reference", token_ids=prefill_ids, shape=list(values.shape),
            row_sha256=[hashlib.sha256(row.tobytes()).hexdigest() for row in values[0]],
            argmax=[int(row.argmax()) for row in values[0]], load_seconds=loaded-started,
            forward_seconds=finished-loaded, mtp=mtp, prompt=prompt, power_percent=100)),flush=True)
        if mtp:
            def fingerprint(array):
                values = np.asarray(array.astype(mx.float32)).astype('<f4')
                assert np.isfinite(values).all()
                return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())

            def record(op, logits=None, hidden=None, **fields):
                mx.eval(*[a for a in (logits, hidden) if a is not None])
                print(json.dumps(dict(event="mtp_reference", op=op,
                    logits=None if logits is None else fingerprint(logits),
                    hidden=None if hidden is None else fingerprint(hidden), **fields)), flush=True)

            body_hidden = model.language_model.model._last_widened
            out, tail_hidden = model(mx.array([[ids[-1]]]), cache=cache, return_hidden=True)
            record("trunk", out, tail_hidden, token_ids=[ids[-1]])
            phase = "decode"
            head_cache = model.make_mtp_cache()
            history = model.mtp_update_cache(body_hidden, mx.array([ids[1:]]), mtp_cache=head_cache)
            record("history", hidden=history, token_ids=ids[1:], offset=head_cache[0].offset)
            base_offset = head_cache[0].offset
            first = int(mx.argmax(out[0, -1]).item())
            token, draft_hidden = first, tail_hidden
            drafts = []
            for depth in range(3):
                out, draft_hidden = model.mtp_forward(draft_hidden, mx.array([[token]]),
                    mtp_cache=head_cache, return_hidden=True)
                record("draft", out, draft_hidden, token_ids=[token], depth=depth,
                    offset=head_cache[0].offset)
                token = int(mx.argmax(out[0, -1]).item())
                drafts.append(token)
            window = [first, *drafts]
            snapshots = [None if c.is_trimmable() else list(c.state) for c in cache]
            with model.verify_capture_scope():
                out, verified_hidden = model(mx.array([window]), cache=cache, return_hidden=True)
            record("verify", out, verified_hidden, token_ids=window)
            choices = np.asarray(mx.argmax(out[0], axis=-1)).tolist()
            keep = 1
            for draft, target in zip(drafts, choices):
                if draft != target:
                    break
                keep += 1
            committed = model.commit_verified_window(cache, snapshots,
                keep_tokens=keep, verified_tokens=len(window))
            assert committed
            # generate_mtpk retains the first authoritative draft row and
            # replaces only recursively predicted rows after it.
            head_cache[0].trim(head_cache[0].offset - (base_offset + 1))
            record("commit", keep=keep, verified=len(window), offset=head_cache[0].offset)
            history = None
            if keep > 1:
                history = model.mtp_update_cache(verified_hidden[:, :keep-1], mx.array([window[1:keep]]), mtp_cache=head_cache)
            record("history_commit", hidden=history, token_ids=window[1:keep], offset=head_cache[0].offset)
            token = int(choices[keep-1])
            out, hidden = model(mx.array([[token]]), cache=cache, return_hidden=True)
            record("continuation", out, hidden, token_ids=[token])
            next_token = int(mx.argmax(out[0, -1]).item())
            out, hidden = model.mtp_forward(hidden, mx.array([[next_token]]), mtp_cache=head_cache, return_hidden=True)
            record("draft_continuation", out, hidden, token_ids=[next_token], offset=head_cache[0].offset)
            print(json.dumps(dict(event="mtp_reference_end", keep=keep, verified=len(window))), flush=True)
        if decode:
            # Numerical cache-continuation oracle, not the serving sampler or
            # a throughput benchmark. Follow greedy tokens to natural EOS.
            phase = "decode"
            token = int(values[0, -1].argmax())
            generated = []
            while token not in tokenizer.eos_token_ids:
                generated.append(token)
                started = time.monotonic()
                out = model(mx.array([[token]], dtype=mx.int32), cache=cache)
                mx.eval(out)
                elapsed = time.monotonic() - started
                values = np.asarray(out.astype(mx.float32)).astype('<f4')
                assert np.isfinite(values).all()
                next_token = int(values[0, -1].argmax())
                print(json.dumps(dict(event="decode_reference", step=len(generated), token_id=token,
                    shape=list(values.shape), argmax=next_token,
                    row_sha256=hashlib.sha256(values[0, -1].tobytes()).hexdigest(),
                    forward_seconds=elapsed)), flush=True)
                token = next_token
            print(json.dumps(dict(event="decode_reference_end", eos_token=token,
                eos_token_ids=sorted(tokenizer.eos_token_ids), token_ids=generated,
                text=tokenizer.decode(generated), finish_reason="eos")), flush=True)
        phase = "finishing"
        del out, model, cache
    finally:
        stop.set()
        if canary:
            prober.join(timeout=2)
            watcher.join(timeout=2)
        print(json.dumps(dict(event="reference_canary_summary", enabled=canary, samples=samples,
            incomplete=prober.is_alive())),file=sys.stderr,flush=True)
        if prober.is_alive():
            os._exit(86)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("root")
    parser.add_argument("bridge")
    parser.add_argument("--decode", action="store_true")
    parser.add_argument("--server-defaults", action="store_true")
    parser.add_argument("--mtp", action="store_true", help="Numerical draft/verify/commit round, not a generation benchmark")
    parser.add_argument("--ar", action="store_true", help="Original cold prefill and sampled AR to natural EOS, with final-state capture")
    parser.add_argument("--mtp-prefill", action="store_true", help="Original streaming committed-history prefill and first draft")
    parser.add_argument("--mtp-generate", action="store_true", help="Original sampled generate_mtpk to natural EOS")
    parser.add_argument("--readme", type=Path, help="Summarize this README instead of ping; freezes the exact prompt in the receipt")
    parser.add_argument("--restored-suffixes", type=Path, help="After MTP EOS, compare live-cache suffix extension using short text, this README, then short Python instructions")
    parser.add_argument("--snapshot-restore", action="store_true", help="Snapshot and restore fresh trunk/MTP containers before each suffix; verify retained snapshots after mutation")
    parser.add_argument("--chat", action="store_true", help="Full README-summary/story/Python chat through the original SessionBank; needs --readme and --ar or --mtp-generate")
    parser.add_argument("--source-run", type=Path, help="Match a completed warm model-eval chat: settings, warmup, bootstrap, output and cache frontiers")
    parser.add_argument("--canary", choices=("on", "off"), default="on")
    parser.add_argument("--token-events", choices=("full", "progress"), default="full",
        help="Chat token trace or at most one real decode-progress event per second; full output is always retained")
    args = parser.parse_args()
    assert sum((args.mtp, args.decode, args.ar, args.mtp_prefill, args.mtp_generate)) <= 1, "choose exactly one execution oracle"
    assert not args.restored_suffixes or args.mtp_generate, "restored suffix oracle requires --mtp-generate"
    assert not args.snapshot_restore or args.restored_suffixes, "snapshot restore requires suffix cases"
    assert not args.chat or (args.readme and (args.ar or args.mtp_generate) and not args.restored_suffixes), "chat requires README and a generation lane, without suffix-only cases"
    assert not args.source_run or (args.chat and args.server_defaults), "source matching requires chat and server defaults"
    assert args.token_events == "full" or args.chat, "progress-only token events require chat"
    source = source_chat([json.loads(line) for line in args.source_run.read_text().splitlines()
        if line.startswith("{")], args.mtp_generate, args.canary == "on") if args.source_run else None
    main(args.root, args.bridge, args.decode, args.server_defaults, args.mtp, args.ar, args.readme, args.mtp_prefill, args.mtp_generate, args.restored_suffixes, args.snapshot_restore, args.chat, source, args.canary == "on", args.token_events)