"""Load raw per-attempt records into a tidy DataFrame with derived columns (scenario costs, method labels)."""
from __future__ import annotations
import glob, json
import pandas as pd
import numpy as np
from ..util import ROOT, read_jsonl

TEXT_METHODS = ["A0", "A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8"]
CACHE_METHODS = ["B0", "B1", "B2", "B2b", "B3", "B4", "B5", "B6", "B9", "C0", "C2", "C3"]
LABELS = {"A0": "A0 no message", "A1": "A1 full text", "A2": "A2 head/tail", "A3": "A3 extractive", "A4": "A4 summary", "A5": "A5 structured evidence", "A6": "A6 refs+retrieval",
          "A7": "A7 zstd text", "A8": "A8 LLMLingua-2", "B0": "B0 native cache", "B1": "B1 serialized reuse", "B2": "B2 packed quant", "B3": "B3 sink+recent", "B4": "B4 random",
          "B5": "B5 attention (SnapKV-style)", "B6": "B6 pyramid", "B9": "B9 select+quant", "C0": "C0 identity (neg. control)", "C2": "C2 linear mapper", "C3": "C3 mapper+select"}

def load_experiments(patterns: list[str], last_attempt_only: bool = True) -> pd.DataFrame:
    """Load records. By default keep the LAST attempt per trial (retries supersede); all attempts are kept when
    last_attempt_only=False (used for the failure/attempt tables)."""
    rows = []
    for pat in patterns:
        for p in glob.glob(str(ROOT / "runs/raw" / pat)):
            for r in read_jsonl(p):
                rows.append(flatten(r))
    df = pd.DataFrame(rows)
    if last_attempt_only and len(df):
        df = df.sort_values("timestamp_utc").drop_duplicates("trial_id", keep="last")
    return df

def _scenario_costs(r: dict) -> dict:
    sp = r.get("timing_spans") or {}
    m = r["method_id"]
    enc = sp.get("encode_ms", 0.0) or 0.0
    recv = (sp.get("receiver_prefill_ms", 0.0) or 0.0) + (sp.get("ttft_ms", 0.0) or 0.0) + (sp.get("decode_ms", 0.0) or 0.0)
    if m.startswith("B") or m.startswith("C"):
        doc_prefill = sp.get("doc_prefill_ms", 0.0) or 0.0
        compress = (sp.get("compress_ms", 0.0) or 0.0) + (sp.get("serialize_ms", 0.0) or 0.0)
        load = sp.get("load_ms", 0.0) or 0.0
        # cold: document must be read & cached, compressed, loaded, then query processed
        cold = doc_prefill + compress + load + recv
        # warm (cache already built & compressed once, reused across N queries): per-query = load + query
        warm_query = load + recv
        reuse_n = max(1, r.get("reuse_count") or 1)
        amortized = (doc_prefill + compress) / reuse_n + warm_query
        return {"cost_cold_ms": cold, "cost_warm_query_ms": warm_query, "cost_amortized_ms": amortized, "cost_build_ms": doc_prefill + compress, "receiver_ms": recv}
    else:
        reuse_n = max(1, r.get("reuse_count") or 1)
        cold = enc + recv
        warm_query = recv  # message already built
        amortized = enc / reuse_n + recv
        return {"cost_cold_ms": cold, "cost_warm_query_ms": warm_query, "cost_amortized_ms": amortized, "cost_build_ms": enc, "receiver_ms": recv}

def flatten(r: dict) -> dict:
    sc = r.get("task_scores") or {}
    cc = r.get("compression_config") or {}
    out = {k: r.get(k) for k in ["experiment_id", "trial_id", "attempt_id", "timestamp_utc", "run_group_id", "task_family", "document_id", "question_id", "split", "cluster_id", "method_id",
                                 "question_access", "topology", "handoff_index", "hardware_id", "reuse_count", "status", "error_code", "serialized_payload_bytes", "message_utf8_bytes",
                                 "cache_tensor_bytes", "fetched_bytes", "total_wall_ms", "ttft_ms", "peak_prefill_memory", "peak_decode_memory", "prediction_artifact", "protocol_hash"]}
    out["sender"] = (r.get("sender_model_revision") or "").split("@")[0].split("/")[-1]
    out["receiver"] = (r.get("receiver_model_revision") or "").split("@")[0].split("/")[-1]
    out["pair"] = f"{out['sender']}->{out['receiver']}" if out["sender"] else out["receiver"]
    out["correct"] = sc.get("correct", np.nan); out["f1"] = sc.get("f1", np.nan); out["em"] = sc.get("em", np.nan); out["pass_rate"] = sc.get("pass_rate", np.nan); out["abstained"] = sc.get("abstained", np.nan)
    out["quality"] = sc.get("f1") if r["task_family"].startswith("public_") else sc.get("correct")
    if r["status"] != "success":
        out["quality_operational"] = 0.0
    else:
        out["quality_operational"] = out["quality"]
    out["budget_tokens"] = cc.get("budget_tokens"); out["bits"] = cc.get("bits"); out["keep_frac"] = cc.get("keep_frac"); out["codec"] = cc.get("codec")
    out["condition"] = r["run_group_id"].split(":", 1)[1] if ":" in (r.get("run_group_id") or "") else r["method_id"]
    out["label"] = LABELS.get(r["method_id"], r["method_id"])
    out["qtype"] = r["question_id"].split("_", 1)[1] if "_" in r["question_id"] else r["task_family"]
    ev = r.get("evidence_artifact") or {}
    if isinstance(ev, dict) and "ledger" in ev:
        L = ev["ledger"]; out["retention_rate"] = L["retention_rate"]; out["unsupported"] = L["unsupported_numbers"] + L["unsupported_identifiers"]; out["distorted"] = L["distorted_sentences"]; out["msg_bytes_hop"] = L["message_bytes"]
    out.update(_scenario_costs(r))
    pp, pdm = r.get("peak_prefill_memory"), r.get("peak_decode_memory")
    out["peak_memory_cold"] = max([x for x in (pp, pdm) if x is not None], default=None)   # cache methods: include the document-build peak
    sp = r.get("timing_spans") or {}
    for k in ["encode_ms", "doc_prefill_ms", "compress_ms", "serialize_ms", "load_ms", "receiver_prefill_ms", "query_prefill_ms", "decode_ms", "sender_prefill_ms"]:
        out["span_" + k] = sp.get(k)
    cost = r.get("cost_with_source") or {}
    out["cost_usd"] = cost.get("usd") if isinstance(cost, dict) else None
    return out