"""Generation eval for Brain-1/Brain-2 checkpoints (plans/09-eval-harness-spec.md).

Generates greedily on `{DATA_DIR}/{dataset}/test.parquet` and scores structured
match. `--baseline` runs a second model on the same set and prints the delta
table (the SFT-2 forgetting check). Decision log: plans/09-eval-decisions.md.

  uv run --group train modal run train/eval.py --limit 64          # smoke
  uv run --group train modal run train/eval.py                     # full SFT-1
  uv run --group train modal run train/eval.py --dataset ha_actions \
      --adapter /checkpoints/brain1/<sft2> --baseline /checkpoints/brain1/ha_actions
"""

import json
import re
import string
from collections import Counter

from modal_common import (
    CHATML,  # train == serve: identical template
    CKPT_DIR,
    DATA_DIR,
    GPU_BRAIN1,
    GPU_BRAIN2,
    GPU_DEFAULTS,
    VOLUMES,
    app,
    cfg,
    rl_image,
    train_image,
)

# Stock transformers+peft, NOT unsloth: its fast-inference kernel breaks on
# left-padded batched generate (RoPE broadcast error). Eval has no speed need.
with train_image.imports():
    import torch
    from datasets import load_dataset
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

KNOWN_KEYS = ("intents", "response", "automation")


# ---------- scoring (pure python: runs anywhere, unit-testable locally) ----------

def parse_channel(text: str):
    """-> (channel, payload). channel=None if not a JSON object with exactly
    one known top key; 'yaml' channel is decided by the GOLD side, not here."""
    try:
        obj = json.loads(text)
    except (json.JSONDecodeError, TypeError):
        return None, None
    if isinstance(obj, dict) and len(obj) == 1 and next(iter(obj)) in KNOWN_KEYS:
        key = next(iter(obj))
        return key, obj[key]
    return None, None


def norm_slot(v):
    """'50', '50%', 50 and 50.0 all compare equal; strings casefold."""
    if isinstance(v, bool):
        return v
    if isinstance(v, (int, float)):
        return float(v)
    if isinstance(v, str):
        s = v.strip().lower()
        try:
            return float(s.removesuffix("%").strip())
        except ValueError:
            return s
    return v


def call_sig(intent: dict):
    slots = intent.get("slots") or {}
    if not isinstance(slots, dict):
        slots = {}
    return (intent.get("name"), tuple(sorted((k, repr(norm_slot(v))) for k, v in slots.items())))


def norm_text(s: str) -> str:
    """SQuAD-style: casefold, drop punctuation + articles, collapse spaces."""
    s = s.lower().translate(str.maketrans("", "", string.punctuation))
    s = re.sub(r"\b(a|an|the)\b", " ", s)
    return " ".join(s.split())


def token_f1(gold: str, pred: str) -> float:
    g, p = norm_text(gold).split(), norm_text(pred).split()
    common = sum((Counter(g) & Counter(p)).values())
    if not common:
        return 0.0
    prec, rec = common / len(p), common / len(g)
    return 2 * prec * rec / (prec + rec)


def _intent_list(payload):
    if not isinstance(payload, list):
        return None
    if not all(isinstance(i, dict) and i.get("name") for i in payload):
        return None
    return payload


def slices_of(gold_intents: list[dict]) -> list[str]:
    """Failure clusters called out in the spec, keyed on the GOLD call."""
    out = []
    if len(gold_intents) > 1:
        out.append("multi_intent")
    for it in gold_intents:
        name, slots = it.get("name", ""), it.get("slots") or {}
        target = str(slots.get("name", "")).lower()
        if name in ("HassTurnOn", "HassTurnOff") and (
            "lock" in target or slots.get("domain") == "lock"
        ):
            out.append("locks")
        if "Timer" in name or "Vacuum" in name:
            out.append("timer_vacuum")
        if name in (
            "HassSetPosition",
            "HassLightSet",
            "HassClimateSetTemperature",
            "HassSetVolume",
            "HassHumidifierSetpoint",
            "HassHumidifierMode",
        ):
            out.append("attr_calls")
    return sorted(set(out))


def score_row(gold_text: str, pred_text: str, sys_prompt: str = "") -> dict:
    gold_ch, gold_pl = parse_channel(gold_text)
    if gold_ch is None:  # gold isn't brain-1 JSON -> Brain-2 YAML row
        return _score_yaml(gold_text, pred_text)

    pred_ch, pred_pl = parse_channel(pred_text)
    r = {
        "bucket": gold_ch,
        "json_valid": pred_ch is not None,
        "route_ok": pred_ch == gold_ch,
        "slices": [],
    }
    if gold_ch == "intents":
        gold_il = _intent_list(gold_pl) or []
        r["slices"] = slices_of(gold_il)
        # acon96 noise: ~10% of golds name entities absent from the prompt's
        # device list (corrupted/derived ids) — unwinnable rows; report split out.
        if sys_prompt:
            names = [str((i.get("slots") or {}).get("name", "")) for i in gold_il]
            r["gold_grounded"] = all((not n) or n in sys_prompt for n in names)
        pred_il = _intent_list(pred_pl) if pred_ch == "intents" else None
        if pred_il is None:
            r.update(name_match=False, full_set=False, full_ord=False)
        else:
            gs, ps = [call_sig(i) for i in gold_il], [call_sig(i) for i in pred_il]
            r["name_match"] = Counter(i["name"] for i in gold_il) == Counter(
                i["name"] for i in pred_il
            )
            r["full_set"] = sorted(map(repr, gs)) == sorted(map(repr, ps))
            r["full_ord"] = gs == ps
    elif gold_ch == "response":
        p = pred_pl if (pred_ch == "response" and isinstance(pred_pl, str)) else ""
        r["exact"] = bool(p) and norm_text(gold_pl) == norm_text(p)
        r["f1"] = token_f1(gold_pl, p) if p else 0.0
    elif gold_ch == "automation":
        p = pred_pl if (pred_ch == "automation" and isinstance(pred_pl, str)) else ""
        r["verbatim"] = p == gold_pl  # contract: hand off VERBATIM
        r["norm_match"] = bool(p) and norm_text(gold_pl) == norm_text(p)
    return r


def _score_yaml(gold_text: str, pred_text: str) -> dict:
    import yaml

    def load(s):
        try:
            doc = yaml.safe_load(s)
            return doc if isinstance(doc, (dict, list)) else None
        except yaml.YAMLError:
            return None

    gold_doc, pred_doc = load(gold_text), load(pred_text)
    return {
        "bucket": "yaml",
        "json_valid": pred_doc is not None,  # "valid output" for brain-2 = parseable YAML
        "route_ok": pred_doc is not None,
        "slices": [],
        "struct_match": gold_doc is not None and gold_doc == pred_doc,
        "exact": gold_text.strip() == pred_text.strip(),
    }


def _grounded_split(action_rows, flag, pct):
    sel = [r for r in action_rows if r.get("gold_grounded") is flag]
    return {"n": len(sel), "full_set": pct(sel, "full_set")} if sel else None


def aggregate(rows: list[dict]) -> dict:
    def pct(items, key):
        items = [r for r in items if key in r]
        return round(100 * sum(r[key] for r in items) / len(items), 2) if items else None

    by = {b: [r for r in rows if r["bucket"] == b] for b in (*KNOWN_KEYS, "yaml")}
    m = {
        "rows": len(rows),
        "json_valid": pct(rows, "json_valid"),
        "routing": pct(rows, "route_ok"),
        "buckets": {b: len(v) for b, v in by.items() if v},
        "action": {
            "n": len(by["intents"]),
            "name_match": pct(by["intents"], "name_match"),
            "full_set": pct(by["intents"], "full_set"),
            "full_ord": pct(by["intents"], "full_ord"),
            "grounded": _grounded_split(by["intents"], True, pct),
            "ungrounded": _grounded_split(by["intents"], False, pct),
        },
        "response": {
            "n": len(by["response"]),
            "exact": pct(by["response"], "exact"),
            "f1": round(
                sum(r["f1"] for r in by["response"]) / len(by["response"]), 3
            )
            if by["response"]
            else None,
        },
        "automation": {
            "n": len(by["automation"]),
            "verbatim": pct(by["automation"], "verbatim"),
            "norm_match": pct(by["automation"], "norm_match"),
        },
        "yaml": {
            "n": len(by["yaml"]),
            "valid": pct(by["yaml"], "json_valid"),
            "struct_match": pct(by["yaml"], "struct_match"),
        },
        "slices": {},
    }
    for s in ("multi_intent", "locks", "timer_vacuum", "attr_calls"):
        sl = [r for r in by["intents"] if s in r["slices"]]
        if sl:
            m["slices"][s] = {"n": len(sl), "full_set": pct(sl, "full_set")}
    return m


def report(m: dict, title: str) -> str:
    L = [f"==== {title} | {m['rows']} rows ===="]
    L.append(f"json-valid {m['json_valid']}%   routing {m['routing']}%")
    a = m["action"]
    if a["n"]:
        L.append(
            f"ACTION     ({a['n']:>5}): names {a['name_match']}% | "
            f"full-call set {a['full_set']}% | ordered {a['full_ord']}%"
        )
        if a.get("grounded"):
            g, u = a["grounded"], a.get("ungrounded") or {"n": 0, "full_set": "-"}
            L.append(
                f"  gold GROUNDED ({g['n']:>4}): full-call {g['full_set']}%   "
                f"| corrupted-gold ({u['n']:>4}): {u['full_set']}%  <- headline is GROUNDED"
            )
    r = m["response"]
    if r["n"]:
        L.append(f"RESPONSE   ({r['n']:>5}): exact {r['exact']}% | token-F1 {r['f1']}")
    au = m["automation"]
    if au["n"]:
        L.append(
            f"AUTOMATION ({au['n']:>5}): verbatim {au['verbatim']}% | "
            f"normalized {au['norm_match']}%"
        )
    y = m["yaml"]
    if y["n"]:
        L.append(f"YAML       ({y['n']:>5}): valid {y['valid']}% | struct {y['struct_match']}%")
    for s, v in m["slices"].items():
        L.append(f"  slice {s:<13} ({v['n']:>4}): full-call set {v['full_set']}%")
    return "\n".join(L)


def delta_report(base: dict, new: dict) -> str:
    """Forgetting table: new minus baseline on the same test set (BWT-style,
    Lopez-Paz & Ranzato 2017). Negative delta on old-skill metrics = forgetting."""
    pairs = [
        ("json-valid %", ("json_valid",)),
        ("routing %", ("routing",)),
        ("action names %", ("action", "name_match")),
        ("action full-call set %", ("action", "full_set")),
        ("action full-call GROUNDED %", ("action", "grounded", "full_set")),
        ("action full-call ordered %", ("action", "full_ord")),
        ("response exact %", ("response", "exact")),
        ("automation verbatim %", ("automation", "verbatim")),
    ]
    L = [f"{'metric':<28}{'baseline':>10}{'new':>10}{'delta':>10}"]
    for label, path in pairs:
        b, n = base, new
        for k in path:
            b, n = b.get(k) if b else None, n.get(k) if n else None
        if b is None and n is None:
            continue
        d = round(n - b, 2) if (b is not None and n is not None) else None
        flag = "  <- forgetting" if (d is not None and d < -1) else ""
        L.append(f"{label:<28}{b!s:>10}{n!s:>10}{d!s:>10}{flag}")
    return "\n".join(L)


# ---------- generation (GPU container only) ----------

def _eval(
    dataset: str, adapter: str, base_model: str, limit: int, batch: int, max_new: int,
    split: str = "test",
) -> dict:
    src = adapter or base_model
    print(f"[eval] cuda={torch.cuda.is_available()} src={src} dataset={dataset}")
    tok = AutoTokenizer.from_pretrained(src, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        base_model, torch_dtype=torch.bfloat16, trust_remote_code=True
    ).to("cuda")
    if adapter:
        model = PeftModel.from_pretrained(model, adapter)
    model.eval()
    # brain2's llama-format vocab has no <|im_end|>; its training closed turns
    # with the real EOS (train_grpo._load) — mirror whichever this vocab has.
    turn_end = "<|im_end|>" if "<|im_end|>" in tok.get_vocab() else tok.eos_token
    tok.chat_template = CHATML.replace("<|im_end|>", turn_end)
    tok.padding_side = "left"  # decoder-only batching: pad left or completions break
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token or turn_end
    end_id = tok.convert_tokens_to_ids(turn_end)
    stop_ids = [end_id] + ([tok.eos_token_id] if tok.eos_token_id not in (None, end_id) else [])

    # split="train" is a MEMORIZATION check (e.g. yaml_sft has no test split yet) —
    # the report title carries the split so the number can't masquerade as held-out.
    ds = load_dataset("parquet", data_files=f"{DATA_DIR}/{dataset}/{split}.parquet", split="train")
    if limit:
        ds = ds.select(range(min(limit, len(ds))))
    # Two row schemas exist: brain-1 sets use `messages` [sys,user,assistant];
    # brain-2 yaml_sft uses TRL prompt/completion ([sys,user] / [assistant]).
    if "messages" in ds.column_names:
        prompt_msgs = [m[:-1] for m in ds["messages"]]
        golds = [m[-1]["content"] for m in ds["messages"]]
    else:
        prompt_msgs = list(ds["prompt"])
        golds = [c[-1]["content"] for c in ds["completion"]]
    sys_prompts = [p[0]["content"] for p in prompt_msgs]
    users = [p[-1]["content"] for p in prompt_msgs]
    # Same render path as training: template -> string -> tokenizer (the
    # tokenize=True path mis-tokenizes on MiniCPM, see sft-trainer-built).
    prompts = [
        tok.apply_chat_template(p, tokenize=False, add_generation_prompt=True)
        for p in prompt_msgs
    ]

    preds = []
    for i in range(0, len(prompts), batch):
        enc = tok(prompts[i : i + batch], return_tensors="pt", padding=True).to("cuda")
        with torch.inference_mode():
            out = model.generate(
                **enc,
                max_new_tokens=max_new,
                do_sample=False,
                eos_token_id=stop_ids,
                pad_token_id=tok.pad_token_id,
            )
        texts = tok.batch_decode(out[:, enc["input_ids"].shape[1] :], skip_special_tokens=False)
        preds += [t.split(turn_end)[0].strip() for t in texts]
        print(f"[eval] {len(preds)}/{len(prompts)}")

    scored = [
        score_row(g, p, sys_prompt=s) for g, p, s in zip(golds, preds, sys_prompts)
    ]
    metrics = aggregate(scored)
    tag = adapter.strip("/").replace("/", "_") if adapter else f"base_{base_model.split('/')[-1]}"
    title = f"{dataset}/{split} @ {adapter or base_model} (raw greedy)"
    print("\n" + report(metrics, title))

    fails = [
        {"i": i, "user": users[i], "gold": golds[i], "pred": preds[i]}
        for i, r in enumerate(scored)
        if not (r.get("route_ok") and r.get("full_set", r.get("exact", r.get("verbatim", True))))
    ]
    print(f"\n[eval] {len(fails)} imperfect rows; first 8:")
    for f in fails[:8]:
        print(f"  user: {f['user']}\n  gold: {f['gold']}\n  pred: {f['pred']}\n  --")

    out_dir = adapter or f"{CKPT_DIR}/evals"
    import pathlib

    pathlib.Path(out_dir).mkdir(parents=True, exist_ok=True)
    # Base-model runs share /checkpoints/evals, so key the file on the model tag too.
    stem = f"eval_{dataset}" if adapter else f"{tag}_{dataset}"
    if split != "test":
        stem += f"_{split}"
    out_file = f"{out_dir}/{stem}{f'_limit{limit}' if limit else ''}.json"
    rows_dump = [
        {"i": i, "user": users[i], "gold": g, "pred": p, **r}
        for i, (g, p, r) in enumerate(zip(golds, preds, scored))
    ]
    with open(out_file, "w") as fh:
        json.dump({"meta": {"dataset": dataset, "split": split, "adapter": adapter,
                            "base": base_model, "limit": limit, "mode": "raw_greedy", "tag": tag},
                   "metrics": metrics, "rows": rows_dump}, fh, ensure_ascii=False, indent=1)
    VOLUMES[CKPT_DIR].commit()
    print(f"[eval] per-row results -> {out_file}")
    return metrics


# Eval is deterministic + cheap to re-run: retries would only mask real bugs.
_EVAL_OPTS = {**GPU_DEFAULTS, "retries": 0}
# brain2 evals run on rl_image: its pinned transformers is the only stack that
# loads the llama-format brain2 repo (no model_type in config), and it's the
# stack brain2 trains/serves on — train == serve applies to the image too.
_EVAL_OPTS_B2 = {**_EVAL_OPTS, "image": rl_image}


@app.function(gpu=GPU_BRAIN1, **_EVAL_OPTS)
def eval_brain1(dataset: str, adapter: str, base_model: str, limit: int, batch: int, max_new: int, split: str):
    return _eval(dataset, adapter, base_model, limit, batch, max_new, split)


@app.function(gpu=GPU_BRAIN2, **_EVAL_OPTS_B2)
def eval_brain2(dataset: str, adapter: str, base_model: str, limit: int, batch: int, max_new: int, split: str):
    return _eval(dataset, adapter, base_model, limit, batch, max_new, split)


@app.local_entrypoint()
def main(
    dataset: str = "ha_actions",
    adapter: str = f"{CKPT_DIR}/brain1/ha_actions",
    baseline: str = "",  # second model on the same set; "base" = bare base model
    brain: str = "brain1",
    limit: int = 0,
    batch: int = 16,
    max_new: int = 256,
    split: str = "test",
):
    fn = {"brain1": eval_brain1, "brain2": eval_brain2}[brain]
    base_model = cfg["base_models"][brain]
    new = fn.remote(dataset, "" if adapter == "base" else adapter, base_model, limit, batch, max_new, split)
    if baseline:
        old = fn.remote(
            dataset, "" if baseline == "base" else baseline, base_model, limit, batch, max_new, split
        )
        print(f"\n==== forgetting check on {dataset}: {baseline} -> {adapter} ====")
        print(delta_report(old, new))