"""Trial planner and executor for single-handoff experiments (Tracks A and B; one process; laptop).

A *condition* is a method + parameters + question-access regime. A *trial* is (condition, model pair, document, question).
Question-blind messages/caches are built once per (condition, document) and reused across that document's questions —
this is the reuse scenario, and the build cost is recorded on every trial as `doc_build_ms` with `reuse_count` = number of
questions that shared it. The runner is resumable: successful trial keys are skipped; failures are kept and retried once.
"""
from __future__ import annotations
import json, time, traceback, yaml, sys, os
from pathlib import Path
import numpy as np
import mlx.core as mx

from .util import ROOT, read_json, write_json, sha256_json, sha256_text, utc_now, Timer, read_jsonl
from .records import new_record, write_record, trial_key
from .data.prepare import load_doc, load_key
from .scoring import score_synthetic, score_public
from .data.exec_tasks import run_hidden_tests
from .prompts import SYSTEM_RECEIVER, SYSTEM_RECEIVER_NOMSG, receiver_user_text, receiver_question_only
from .backends.mlx_backend import MLXModel, deep_copy_cache, memory_snapshot, reset_peak
from .cache.sparse import quantize_cache, cache_tensor_bytes, to_sparse_all
from .cache.selection import select_cache
from .cache.serialize import build_manifest, serialize_cache, deserialize_cache, codec_compress, codec_decompress
from .methods import text as TM

FIDELITY = {"C0": "n/a", "C2": "custom", "C3": "custom", "A0": "custom", "A1": "custom", "A2": "custom", "A3": "custom", "A4": "custom", "A5": "custom", "A6": "custom", "A7": "custom",
            "A8": "official", "B0": "custom", "B1": "custom", "B2": "custom", "B2b": "inspired", "B3": "inspired", "B4": "custom",
            "B5": "faithful_reimplementation", "B6": "faithful_reimplementation", "B9": "custom"}

import platform as _pf
HARDWARE_ID = {"MacBookPro-de-simon-pierre-2.local": "laptop-m5max-48g", "MacStudnpierre2": "m3u96a-m3ultra-96g", "m3u96b": "m3u96b-m3ultra-96g"}.get(_pf.node().split(".")[0] if _pf.node().startswith("MacStud") or _pf.node().startswith("m3u") else _pf.node(), _pf.node())

def _load_models(cfg):
    mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]
    cache = {}
    def get(alias):
        if alias not in cache:
            cache[alias] = MLXModel(mcfg[alias]["repo"])
        return cache[alias]
    return get, mcfg

def _select_items(cfg):
    man = read_json(ROOT / "data/manifests/splits.json")
    items = [it for it in man["items"] if it["split"] in cfg["splits"] and it["family"] in cfg["families"]]
    if cfg.get("lengths"):
        items = [it for it in items if it.get("target_tokens") is None or it["target_tokens"] in cfg["lengths"]]
    lim = cfg.get("limit_docs_per_family")
    if lim:
        out = []
        for fam in cfg["families"]:
            out += [it for it in items if it["family"] == fam][:lim]
        items = out
    return items

def _cond_id(c: dict) -> str:
    parts = [c["method"]]
    for k in sorted(c):
        if k not in ("method", "access", "note", "families", "max_questions"):
            parts.append(f"{k}={c[k]}")
    if "access" in c:
        parts.append(c["access"])
    return "|".join(parts)

def plan(config_path: str, output: str) -> int:
    cfg = yaml.safe_load(open(config_path))
    items = _select_items(cfg)
    rows = []
    pairs = cfg["model_pairs"]
    for it in items:
        doc = load_doc(it["document_id"], it["family"])
        nq0 = len(doc["questions"]) if not cfg.get("max_questions_per_doc") else min(len(doc["questions"]), cfg["max_questions_per_doc"])
        if cfg.get("question_range"): nq0 = len(doc["questions"][cfg["question_range"][0]:cfg["question_range"][1]])
        for pair in pairs:
            for c in cfg["conditions"]:
                if c.get("families") and it["family"] not in c["families"]:
                    continue
                nq = min(nq0, c["max_questions"]) if c.get("max_questions") else nq0
                for seed in cfg.get("seeds", [0]):
                    rows.append({"experiment_id": cfg["experiment_id"], "condition": _cond_id(c), "method": c["method"], "access": c.get("access", "none"),
                                 "sender": pair["sender"], "receiver": pair["receiver"], "family": it["family"], "document_id": it["document_id"],
                                 "split": it["split"], "target_tokens": it.get("target_tokens"), "n_questions": nq, "seed": seed, "expected_trials": nq})
    import csv
    Path(output).parent.mkdir(parents=True, exist_ok=True)
    with open(output, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys())); w.writeheader(); w.writerows(rows)
    tot = sum(r["expected_trials"] for r in rows)
    print(f"plan: {len(items)} documents, {len(cfg['conditions'])} conditions, {len(pairs)} model pairs -> {len(rows)} doc-conditions, {tot} trials -> {output}")
    return 0

# --------------------------------------------------------------------------------------------------------------------
class DocContext:
    """Per-document receiver state: tokenized prefix and dense cache (built lazily once)."""
    def __init__(self, receiver: MLXModel, doc_text: str, family: str):
        self.receiver = receiver; self.doc_text = doc_text; self.family = family
        self._dense = None; self.build_ms = None; self.prefix_ids = None; self.full_ids_fn = None
        self.msg_cache = {}
        self.sys = SYSTEM_RECEIVER
    def prefix_tokens(self):
        """Receiver prefix = chat tokens up to (and including) the document and the blank line before 'Question:'.
        Computed as the longest common token prefix between the prefix text and a full prompt with a canonical question."""
        if self.prefix_ids is None:
            pre_txt = self.receiver.tokenizer.apply_chat_template([{"role": "system", "content": self.sys}, {"role": "user", "content": "Document:\n" + self.doc_text + "\n\n"}], add_generation_prompt=False, tokenize=False)
            pre = self.receiver.encode(pre_txt)
            full = self.receiver.chat_prompt_tokens([{"role": "system", "content": self.sys}, {"role": "user", "content": receiver_user_text("Document", self.doc_text, "What is stated?", self.family)}])
            n = 0
            while n < min(len(pre), len(full)) and pre[n] == full[n]:
                n += 1
            self.prefix_ids = full[:n]
        return self.prefix_ids
    def split_ids(self, question: str):
        """Full chat tokens for (doc, question); returns (prefix_ids, query_ids) with prefix a true token prefix."""
        full = self.receiver.chat_prompt_tokens([{"role": "system", "content": self.sys}, {"role": "user", "content": receiver_user_text("Document", self.doc_text, question, self.family)}])
        pre = self.prefix_tokens()
        if full[:len(pre)] != pre:
            raise RuntimeError("prefix_boundary: question tokenization altered the document prefix")
        return pre, full[len(pre):]
    def dense_cache(self, prefix_ids):
        if self._dense is None or self._dense_len != len(prefix_ids):
            reset_peak()
            c, t = self.receiver.prefill(prefix_ids)
            self._dense = c; self._dense_len = len(prefix_ids); self.build_ms = t["prefill_ms"]; self.build_peak = int(mx.get_peak_memory())
        return self._dense

def _score(family: str, pred: str, qkey: dict, key: dict) -> dict:
    if family == "exec":
        r = run_hidden_tests(pred, key["hidden_tests"])
        return {"correct": r["correct"], "pass_rate": r["pass_rate"], "error": r.get("error")}
    if family.startswith("public_"):
        return score_public(pred, [qkey["answer"]] + list(qkey.get("aliases", [])))
    return score_synthetic(pred, qkey)

def run(config_path: str, resume: bool = False, limit: int | None = None, dry_run: bool = False) -> int:
    cfg = yaml.safe_load(open(config_path))
    res = yaml.safe_load(open(ROOT / "configs/resources.yaml"))
    protocol_hash = cfg.get("protocol_hash") or "unfrozen"
    config_hash = sha256_json(cfg)
    out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"
    done = set()
    if resume and out_path.exists():
        for r in read_jsonl(out_path):
            if r["status"] == "success":
                done.add(r["trial_id"])
    items = _select_items(cfg)
    get_model, mcfg = _load_models(cfg)
    gen_cfg = cfg.get("generation", {})
    gen_max = gen_cfg.get("max_tokens", 48)
    n_done = n_run = n_fail = 0
    log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log"
    log_path.parent.mkdir(parents=True, exist_ok=True)
    def log(msg):
        line = f"{utc_now()} {msg}"
        print(line, flush=True)
        with open(log_path, "a") as f: f.write(line + "\n")
    log(f"run start experiment={cfg['experiment_id']} config_hash={config_hash[:12]} protocol={protocol_hash[:12]} items={len(items)} resume={resume} already_done={len(done)}")
    if dry_run:
        return plan(config_path, str(ROOT / "protocol" / f"{cfg['experiment_id']}_dryrun.csv"))
    t_start = time.time()
    for it in items:
        fam = it["family"]; doc = load_doc(it["document_id"], fam); key = load_key(it["document_id"], fam)
        qs = doc["questions"][: cfg.get("max_questions_per_doc") or None]
        if cfg.get("question_range"):   # exploratory extensions: a slice of the question list, e.g. [5, 15]
            a, b = cfg["question_range"]; qs = doc["questions"][a:b]
        gen_max = gen_cfg.get("per_family", {}).get(fam, gen_cfg.get("max_tokens", 48))
        qkeys = {q["question_id"]: q for q in key["questions"]}
        for pair in cfg["model_pairs"]:
            receiver = get_model(pair["receiver"]); sender = get_model(pair["sender"])
            dc = DocContext(receiver, doc["text"], fam)
            for cond in cfg["conditions"]:
                if cond.get("families") and fam not in cond["families"]:
                    continue
                cid = _cond_id(cond); method = cond["method"]; access = cond.get("access", "none")
                if method.startswith("B") or method.startswith("C"):
                    access = cond.get("access", "blind")
                for seed in cfg.get("seeds", [0]):
                    # trials for this (doc, cond, pair, seed)
                    pending = []
                    for q in qs[: cond.get("max_questions") or None]:
                        tk = trial_key(cfg.get("trial_key_experiment_id", cfg["experiment_id"]), it["document_id"], q["question_id"], method, {k: v for k, v in cond.items() if k not in ("method", "families", "max_questions")},
                                       pair["sender"], pair["receiver"], access, protocol_hash, seed)
                        if tk in done:
                            n_done += 1; continue
                        pending.append((q, tk))
                    if not pending:
                        continue
                    if limit is not None and n_run >= limit:
                        log(f"limit {limit} reached"); return 0
                    # ---- build the question-blind message / cache once per document (reuse scenario)
                    shared = None; shared_err = None
                    if access in ("blind", "none"):
                        try:
                            shared = _build_message(method, cond, doc["text"], None, sender, receiver, dc, seed)
                        except Exception as e:
                            shared_err = e
                            log(f"BUILD FAIL {cid} {it['document_id']}: {type(e).__name__}: {str(e)[:200]}")
                    for q, tk in pending:
                        rec = _run_one(cfg, cond, cid, method, access, pair, it, doc, q, qkeys.get(q["question_id"]), key, sender, receiver, dc, seed, tk, shared, shared_err, gen_max, protocol_hash, config_hash, len(pending))
                        write_record(rec, out_path)
                        n_run += 1
                        if rec["status"] != "success": n_fail += 1
                        sc = rec["task_scores"] or {}
                        log(f"{cid:>40s} {it['document_id']:>28s} {q['question_id']:>14s} {rec['status']:>10s} corr={sc.get('correct', sc.get('f1'))} bytes={rec['serialized_payload_bytes']} ms={rec['total_wall_ms'] and round(rec['total_wall_ms'])} ans={json.dumps((rec['prediction_artifact'] or '')[:60])}")
                        if limit is not None and n_run >= limit:
                            log(f"limit {limit} reached"); return 0
            mx.clear_cache()
    log(f"run end: ran={n_run} skipped_done={n_done} failed={n_fail} elapsed_s={round(time.time()-t_start)}")
    return 0

def _build_message(method, cond, doc_text, question, sender, receiver, dc: DocContext, seed):
    """Returns dict with message (Track A) or cache payload info (Track B); encode timings included."""
    t = Timer(); t.start("encode")
    out = {"method": method}
    if method == "A0":
        out["msg"] = TM.a0_no_message()
    elif method == "A1":
        out["msg"] = TM.a1_full(doc_text)
    elif method == "A2":
        out["msg"] = TM.a2_truncate(doc_text, cond["budget_tokens"], receiver.tokenizer)
    elif method == "A3":
        out["msg"] = TM.a3_extractive(doc_text, cond["budget_tokens"], receiver.tokenizer, question)
    elif method == "A4":
        out["msg"] = TM.a4_summary(doc_text, cond["budget_tokens"], sender, question)
    elif method == "A5":
        out["msg"] = TM.a5_structured(doc_text, cond["budget_tokens"], sender, question)
        if hasattr(dc, "msg_cache"): dc.msg_cache[("A5", cond["budget_tokens"], question)] = out["msg"]
    elif method == "A6":
        key5 = ("A5", cond["budget_tokens"], question)
        ev = dc.msg_cache.get(key5) if hasattr(dc, "msg_cache") else None
        out["msg"] = TM.a6_references(doc_text, cond["budget_tokens"], sender, question, evidence_msg=ev)
        out["reused_a5"] = ev is not None
    elif method == "A7":
        out["msg"] = TM.a7_codec(doc_text, cond.get("codec", "zstd"))
    elif method == "A8":
        out["msg"] = TM.a8_llmlingua(doc_text, cond["budget_tokens"], receiver.tokenizer, question, device=cond.get("device", "mps"))
    elif method.startswith("C"):
        # Track C: cross-checkpoint transfer. Sender cache built on the receiver's exact prefix token ids (shared tokenizer).
        from .mapper import LinearKVMapper, identity_transfer_negative_control
        pre = dc.prefix_tokens()
        assert sender.info.tokenizer_hash == receiver.info.tokenizer_hash, "C methods need a shared tokenizer"
        t.start("sender_prefill"); s_cache, tp = sender.prefill(pre); t.stop("sender_prefill")
        t.start("compress")
        if method == "C0":
            lm = [int(round(lr * (sender.info.num_layers - 1) / (receiver.info.num_layers - 1))) for lr in range(receiver.info.num_layers)]
            c = identity_transfer_negative_control(s_cache, receiver.info, lm); info = {"negative_control": True}
        else:
            mp = LinearKVMapper.load(ROOT / cond["mapper"], sender.info, receiver.info)
            c = mp.apply(s_cache); info = {"mapper": cond["mapper"]}
            if method == "C3":
                window = dc.split_ids(question)[1] if question is not None else receiver.encode("\n\nSummarize the key facts, identifiers, amounts, dates and statuses of this document.")
                c, sel = select_cache(receiver, c, cond.get("policy", "attention"), cond["keep_frac"], window_token_ids=window, seed=seed, window_keep=cond.get("window_keep", 32)); info.update(sel)
        mx.eval([x for cc in c for x in (cc.keys, cc.values)]); mx.synchronize()
        t.stop("compress")
        t.start("serialize")
        man = build_manifest(receiver.info, pre, c, "doc", method, {k: v for k, v in cond.items() if k != "method"}, {"selection_info": {k: v for k, v in info.items() if k != "layer_budgets"}})
        payload = serialize_cache(c, man)
        t.stop("serialize")
        out.update({"cache": c, "payload": payload, "payload_bytes": len(payload), "raw_payload_bytes": len(payload), "cache_bytes": cache_tensor_bytes(c), "manifest": man,
                    "compress_ms": t.spans["compress"], "serialize_ms": t.spans["serialize"], "codec": "none", "select_info": info, "doc_build_ms": tp["prefill_ms"], "doc_build_peak": int(mx.get_peak_memory()),
                    "prefix_tokens": len(pre), "dense_bytes": cache_tensor_bytes(s_cache), "sender_prefill_ms": tp["prefill_ms"]})
        del s_cache
    elif method.startswith("B"):
        pre = dc.prefix_tokens()
        dense = dc.dense_cache(pre)
        out["doc_build_ms"] = dc.build_ms; out["doc_build_peak"] = dc.build_peak; out["prefix_tokens"] = len(pre)
        out["dense_bytes"] = cache_tensor_bytes(dense)
        t.start("compress")
        if method == "B0":
            c = to_sparse_all(deep_copy_cache(dense)); info = {}
        elif method == "B1":
            c = to_sparse_all(deep_copy_cache(dense)); info = {"serialize_roundtrip": True}
        elif method == "B2":
            c = quantize_cache(dense, cond["bits"], cond.get("group_size", 64)); info = {"bits": cond["bits"]}
        elif method in ("B3", "B4", "B5", "B6", "B9"):
            policy = {"B3": "sink_recent", "B4": "random", "B5": "attention", "B6": "pyramid", "B9": cond.get("policy", "sink_recent")}[method]
            window = None
            if policy in ("attention", "pyramid"):
                if question is not None:
                    window = dc.split_ids(question)[1]          # question-aware: the fresh query tokens are the observation window
                else:
                    window = receiver.encode("\n\nSummarize the key facts, identifiers, amounts, dates and statuses of this document.")  # question-blind generic instruction
            c, info = select_cache(receiver, dense, policy, cond["keep_frac"], window_token_ids=window, seed=seed, sinks=cond.get("sinks", 4), window_keep=cond.get("window_keep", 32))
            if method == "B9":
                c = quantize_cache(c, cond["bits"], cond.get("group_size", 64)); info["bits"] = cond["bits"]
        else:
            raise NotImplementedError(method)
        mx.eval([x for cc in c for x in (cc.keys if isinstance(cc.keys, tuple) else (cc.keys,))])
        mx.synchronize()
        t.stop("compress")
        t.start("serialize")
        man = build_manifest(receiver.info, pre, c, "doc", method, {k: v for k, v in cond.items() if k != "method"}, {"selection_info": {k: v for k, v in info.items() if k != "layer_budgets"}})
        payload = serialize_cache(c, man)
        t.stop("serialize")
        codec = cond.get("codec", "none")
        comp, comp_ms = codec_compress(payload, codec) if codec != "none" else (payload, 0.0)
        out.update({"cache": c, "payload": comp, "payload_bytes": len(comp), "raw_payload_bytes": len(payload), "cache_bytes": cache_tensor_bytes(c), "manifest": man,
                    "compress_ms": t.spans.get("compress", 0.0), "serialize_ms": t.spans["serialize"] + comp_ms, "codec": codec, "select_info": info})
    else:
        raise NotImplementedError(method)
    t.stop("encode")
    out["encode_ms"] = t.spans["encode"]
    return out

def _run_one(cfg, cond, cid, method, access, pair, it, doc, q, qkey, key, sender, receiver, dc, seed, tk, shared, shared_err, gen_max, protocol_hash, config_hash, reuse_count):
    fam = it["family"]
    rec = new_record(experiment_id=cfg["experiment_id"], trial_id=tk, run_group_id=cfg["experiment_id"] + ":" + cid, protocol_hash=protocol_hash, config_hash=config_hash,
                     task_family=fam, document_id=it["document_id"], document_hash=it["document_hash"], question_id=q["question_id"], split=it["split"], cluster_id=it["cluster_id"],
                     sender_model_revision=f"{sender.info.repo_id}@{sender.info.revision}", receiver_model_revision=f"{receiver.info.repo_id}@{receiver.info.revision}",
                     tokenizer_hashes={"sender": sender.info.tokenizer_hash, "receiver": receiver.info.tokenizer_hash},
                     backend_versions={"mlx": mx.__version__, "mlx_lm": receiver.info.mlx_lm_version}, hardware_id=HARDWARE_ID, device_state={"mlx_active_before": int(mx.get_active_memory())},
                     weight_precision="bf16", method_id=method, implementation_fidelity=FIDELITY.get(method, "custom"),
                     compression_config={k: v for k, v in cond.items() if k != "method"}, question_access=access, topology="single_handoff", handoff_index=1 if method != "A0" else 0,
                     reuse_count=reuse_count if access in ("blind", "none") else 1, cache_hit=(access in ("blind", "none") and method.startswith("B")), cache_compatibility_hash=None,
                     random_seed=seed, query_tokens=None, message_tokens=None, message_utf8_bytes=None, cache_tensor_bytes=None, serialized_payload_bytes=None,
                     input_usage={}, output_usage={}, provider_cached_usage=None, process_memory=None, backend_memory=None, peak_prefill_memory=None, peak_decode_memory=None,
                     mapper_memory=None, cost_with_source=None, prediction_artifact=None, evidence_artifact=None, reference_key_id=f"keys/{fam}/{it['document_id']}.json",
                     task_scores=None, status="implementation_error", error_code=None, log_artifact=f"runs/logs/{cfg['experiment_id']}.log", payload_hash=None, output_hash=None)
    rec["source_tokens_sender"] = len(sender.encode(doc["text"])); rec["source_tokens_receiver"] = len(receiver.encode(doc["text"]))
    t_total0 = time.perf_counter()
    try:
        if shared_err is not None:
            raise shared_err
        built = shared if shared is not None else _build_message(method, cond, doc["text"], q["question"], sender, receiver, dc, seed)
        spans = {"encode_ms": built["encode_ms"]}
        reset_peak()
        if method.startswith("A"):
            msg = built["msg"]
            if method == "A6":
                material, rinfo = TM.a6_retrieve(msg, q["question"], cond.get("max_fetch_bytes", 2000), cond.get("max_lookups", 6))
                rec["fetched_bytes"] = rinfo["fetched_bytes"]; rec["evidence_artifact"] = rinfo
            else:
                material = msg.text
            if method == "A0":
                ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER_NOMSG}, {"role": "user", "content": receiver_question_only(q["question"], fam)}])
            else:
                ids = receiver.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text(msg.material_label, material, q["question"], fam)}])
            spans["decode_payload_ms"] = msg.decode_ms
            g = receiver.generate_text(ids, max_tokens=gen_max)
            spans.update({"receiver_prefill_ms": g["prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"]})
            rec.update({"message_tokens": len(receiver.encode(material)), "message_utf8_bytes": msg.utf8_bytes, "serialized_payload_bytes": msg.payload_bytes,
                        "side_information_bytes": msg.side_information_bytes, "query_tokens": len(ids), "input_usage": {"receiver_prompt_tokens": len(ids), **msg.sender_usage},
                        "payload_hash": sha256_text(material)})
            if msg.meta: rec["compression_config"] = {**rec["compression_config"], "realized": {k: v for k, v in msg.meta.items() if k not in ("question_aware",)}}
            if method == "A5": rec["evidence_artifact"] = {"n_facts": msg.meta.get("n_facts"), "unresolved_quotes": msg.meta.get("unresolved_quotes"), "parse_ok": msg.meta.get("parse_ok")}
        else:
            pre, query_ids = dc.split_ids(q["question"])
            assert len(pre) == built["prefix_tokens"], "prefix mismatch"
            c = built["cache"]
            spans.update({"doc_prefill_ms": built["doc_build_ms"], "compress_ms": built["compress_ms"], "serialize_ms": built["serialize_ms"]})
            t = Timer(); t.start("load")
            if method == "B1" or cond.get("roundtrip", False):
                raw, dms = codec_decompress(built["payload"], built["codec"]) if built["codec"] != "none" else (built["payload"], 0.0)
                c_use, man = deserialize_cache(raw, receiver.info, built["manifest"]["prefix_token_hash"])
                rec["cache_compatibility_hash"] = man["prefix_token_hash"][:16]
            else:
                c_use = deep_copy_cache(c)
            t.stop("load"); spans["load_ms"] = t.spans["load"]
            g = receiver.generate_from_cache(c_use, query_ids, max_tokens=gen_max)
            spans.update({"receiver_prefill_ms": 0.0, "query_prefill_ms": g["query_prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"]})
            rec.update({"message_tokens": built["prefix_tokens"], "message_utf8_bytes": len(doc["text"].encode("utf-8")), "serialized_payload_bytes": built["payload_bytes"],
                        "cache_tensor_bytes": built["cache_bytes"], "side_information_bytes": built["payload_bytes"] - built["cache_bytes"] if built["codec"] == "none" else None,
                        "query_tokens": len(query_ids), "input_usage": {"receiver_prefix_tokens": built["prefix_tokens"], "dense_cache_bytes": built["dense_bytes"], "query_tokens": len(query_ids)},
                        "payload_hash": sha256_text(str(len(built["payload"]))) , "peak_prefill_memory": built["doc_build_peak"]})
            rec["compression_config"] = {**rec["compression_config"], "realized": {k: v for k, v in built["select_info"].items() if k not in ("layer_budgets",)}}
        pred = g["text"]
        rec["prediction_artifact"] = pred; rec["output_hash"] = sha256_text(pred); rec["output_usage"] = {"output_tokens": g["output_tokens"]}
        rec["ttft_ms"] = spans["ttft_ms"]; rec["timing_spans"] = spans
        rec["total_wall_ms"] = (time.perf_counter() - t_total0) * 1000
        rec["peak_decode_memory"] = int(mx.get_peak_memory()); rec["backend_memory"] = memory_snapshot(); rec["process_memory"] = rec["backend_memory"]["process_rss_bytes"]
        if qkey is None:
            rec["status"] = "implementation_error"; rec["error_code"] = "missing_key"
        else:
            rec["task_scores"] = _score(fam, pred, qkey, key)
            rec["status"] = "success" if pred.strip() else "invalid_output"
    except Exception as e:
        rec["status"] = "oom" if "memory" in str(e).lower() else "implementation_error"
        rec["error_code"] = f"{type(e).__name__}: {str(e)[:300]}"
        rec["total_wall_ms"] = (time.perf_counter() - t_total0) * 1000
        traceback.print_exc()
    return rec