"""Multi-segment + abstention-aware scorer for V1 schema dumps.

Schema: {"segments": [{"intent","slots","text"}, ...], "abstain_reason": null|str}

Metrics:
  - schema_valid:        payload parses + has segments list
  - segment_count_em:    pred_n_segs == gold_n_segs
  - intent_em (per-segment ordered match)
  - intent_em_set (multi-set match, order-agnostic)
  - slots_em (ci-subset per segment)
  - abstain_rate         pct flagged unknown / abstain
  - precision_when_accepted   intent_em / (n - abstain_count)
  - escalation_rate      from three-way verifier verdicts

Usage:
    uv run python poc/llm-finetune/training/score_v1.py <dump.json> [--by intent|scenario|verdict]
"""
from __future__ import annotations
import argparse
import json
import sys
from collections import defaultdict, Counter
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT / "poc/llm-finetune/training"))
from three_way_verifier import verify, _mock_regex_parse, parse_llm_output


def normalize_slots(slots: dict | None) -> dict:
    if not isinstance(slots, dict):
        return {}
    return {str(k).lower(): str(v).strip().lower() for k, v in slots.items() if v is not None and str(v) != ""}


def slots_subset(pred: dict, gold: dict) -> bool:
    p = normalize_slots(pred)
    g = normalize_slots(gold)
    return all(p.get(k) == v for k, v in g.items())


def segment_intent_match(pred_segs: list, gold_segs: list, ordered: bool = True) -> bool:
    if len(pred_segs) != len(gold_segs):
        return False
    pi = [s.get("intent") for s in pred_segs]
    gi = [s.get("intent") for s in gold_segs]
    if ordered:
        return pi == gi
    return Counter(pi) == Counter(gi)


def segment_slots_match(pred_segs: list, gold_segs: list) -> bool:
    if len(pred_segs) != len(gold_segs):
        return False
    return all(slots_subset(p.get("slots", {}), g.get("slots", {})) for p, g in zip(pred_segs, gold_segs))


def score(rows: list[dict], use_verifier: bool = False) -> dict:
    n = len(rows)
    schema_ok = seg_count_ok = intent_em_ord = intent_em_set = slots_em = 0
    abstain_count = 0
    intent_em_when_accepted = 0
    accepted_count = 0
    verdict_counts = Counter()
    by_intent = defaultdict(lambda: [0, 0, 0])  # n, intent_hits, slots_hits
    failures = []

    for r in rows:
        # Pred parsing
        pred_payload = r.get("parsed")  # eval dump must contain this
        if pred_payload is None and "raw_output" in r:
            pred_payload = parse_llm_output(r["raw_output"])

        gold = r.get("gold")  # {segments: [...], abstain_reason: ...}
        if gold is None:
            # backward-compat: legacy dumps had gold_intent + gold_parameters
            gold = {
                "segments": [{"intent": r.get("gold_intent"), "slots": r.get("gold_parameters", {}), "text": r.get("text", "")}],
                "abstain_reason": None,
            }

        gold_segs = gold.get("segments", [])
        primary_intent = gold_segs[0]["intent"] if gold_segs else "unknown"
        by_intent[primary_intent][0] += 1

        if not isinstance(pred_payload, dict):
            failures.append({"text": r.get("text", "")[:80], "gold": gold, "pred_raw": r.get("raw_output", "")[:200], "reason": "unparseable"})
            continue
        schema_ok += 1

        pred_segs = pred_payload.get("segments") or []
        pred_abstain = pred_payload.get("abstain_reason")

        if pred_abstain or (len(pred_segs) == 1 and pred_segs[0].get("intent") == "unknown"):
            abstain_count += 1
        else:
            accepted_count += 1
            if segment_intent_match(pred_segs, gold_segs, ordered=True):
                intent_em_when_accepted += 1

        if len(pred_segs) == len(gold_segs):
            seg_count_ok += 1
        if segment_intent_match(pred_segs, gold_segs, ordered=True):
            intent_em_ord += 1
            by_intent[primary_intent][1] += 1
        if segment_intent_match(pred_segs, gold_segs, ordered=False):
            intent_em_set += 1
        if segment_slots_match(pred_segs, gold_segs):
            slots_em += 1
            by_intent[primary_intent][2] += 1

        if use_verifier:
            v = verify(json.dumps(pred_payload), r.get("text", ""), regex_parse_fn=_mock_regex_parse)
            verdict_counts[v.verdict] += 1

        # log first 10 failures for inspection
        if (not segment_intent_match(pred_segs, gold_segs, ordered=True) or
            not segment_slots_match(pred_segs, gold_segs)) and len(failures) < 10:
            failures.append({
                "text": r.get("text", "")[:140],
                "gold_segs": gold_segs,
                "pred_segs": pred_segs,
                "pred_abstain": pred_abstain,
            })

    precision_when_accepted = intent_em_when_accepted / max(accepted_count, 1)
    abstain_rate = abstain_count / max(n, 1)

    return {
        "n": n,
        "schema_valid_pct": schema_ok / max(n, 1),
        "segment_count_em_pct": seg_count_ok / max(n, 1),
        "intent_em_ordered_pct": intent_em_ord / max(n, 1),
        "intent_em_set_pct": intent_em_set / max(n, 1),
        "slots_em_pct": slots_em / max(n, 1),
        "abstain_rate": abstain_rate,
        "abstain_count": abstain_count,
        "accepted_count": accepted_count,
        "precision_when_accepted": precision_when_accepted,
        "intent_em_when_accepted": intent_em_when_accepted,
        "verdict_counts": dict(verdict_counts) if verdict_counts else None,
        "by_intent": {k: {"n": v[0], "intent": v[1], "slots": v[2]} for k, v in by_intent.items()},
        "failures": failures,
    }


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("dump", help="Dump JSON from eval (V1 schema, with rows containing parsed/gold)")
    p.add_argument("--use-verifier", action="store_true", help="Run three-way verifier on each row")
    p.add_argument("--top", type=int, default=20)
    p.add_argument("--show-fails", action="store_true")
    args = p.parse_args()

    dump = json.loads(Path(args.dump).read_text())
    rows = dump["rows"]
    print(f"loaded {len(rows)} rows from {args.dump}\n")

    r = score(rows, use_verifier=args.use_verifier)

    print("=== V1 multi-segment + abstention scoring ===")
    print(f"  n:                          {r['n']}")
    print(f"  schema_valid:               {r['schema_valid_pct']:.1%}")
    print(f"  segment_count_em:           {r['segment_count_em_pct']:.1%}")
    print(f"  intent_em (ordered):        {r['intent_em_ordered_pct']:.1%}    (V32 baseline: 87.4%)")
    print(f"  intent_em (set, unordered): {r['intent_em_set_pct']:.1%}")
    print(f"  slots_em:                   {r['slots_em_pct']:.1%}    (V32 baseline: 80.6%)")
    print(f"  abstain_rate:               {r['abstain_rate']:.1%}    ({r['abstain_count']}/{r['n']})")
    print(f"  accepted:                   {r['accepted_count']}")
    print(f"  PRECISION when accepted:    {r['precision_when_accepted']:.1%}    ⭐ (target: 99.5%)")

    if r["verdict_counts"]:
        print(f"\n=== three-way verifier verdicts ===")
        for verdict, count in r["verdict_counts"].items():
            print(f"  {verdict:15s} {count}/{r['n']} = {count/r['n']:.1%}")

    print(f"\n=== per primary-intent (top {args.top}) ===")
    for intent, stats in sorted(r["by_intent"].items(), key=lambda kv: -kv[1]["n"])[:args.top]:
        n_i, i_em, s_em = stats["n"], stats["intent"], stats["slots"]
        print(f"  {intent:30s} n={n_i:4d}  intent={i_em}/{n_i} ({i_em/n_i:.0%})  slots={s_em}/{n_i} ({s_em/n_i:.0%})")

    if args.show_fails:
        print(f"\n=== first {len(r['failures'])} failures ===")
        for f in r["failures"]:
            print(f"\n  text: {f.get('text','')}")
            print(f"  gold: {f.get('gold_segs')}")
            print(f"  pred: {f.get('pred_segs')}")
            if f.get("pred_abstain"):
                print(f"  abstain: {f['pred_abstain']}")


if __name__ == "__main__":
    main()