"""
ab_local_eval.py — A/B the locally-served base vs LoRA-tuned 12B on the holdout.

Generates completions by POSTing each holdout [system,user] prompt to a running
llama-server (localhost:8110/v1), then scores with vn_validate. Run it once per
arm (base server up, then base+--lora up), writing a labelled output file each
time; a final --compare pass prints the side-by-side.

Usage:
    # with the BASE server running:
    python tools/ab_local_eval.py --arm base
    # with the base+LoRA server running:
    python tools/ab_local_eval.py --arm tuned
    # then:
    python tools/ab_local_eval.py --compare
"""
from __future__ import annotations
import argparse
import json
import os
import sys

import requests

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from vn_validate import validate_turn, ngram_overlap  # noqa: E402

REPO = os.path.join(os.path.dirname(__file__), "..")
HOLDOUT = os.path.join(REPO, "runs/vn12b-r32-v1/holdout_generations.jsonl")
OUT_DIR = os.path.join(REPO, "runs/vn12b-r32-v1")
URL = os.environ.get("ARS_FABULA_BASE_URL", "http://localhost:8110/v1")

# Match the app's real serving config for a fair, representative A/B.
GEN = dict(
    temperature=0.7,
    top_p=0.95,
    max_tokens=700,
    frequency_penalty=0.4,
    presence_penalty=0.3,
    seed=12345,
    stream=False,
)


def load_prompts():
    rows = []
    for line in open(HOLDOUT, encoding="utf-8"):
        line = line.strip()
        if not line:
            continue
        rec = json.loads(line)
        msgs = rec["messages"]
        system = next(m["content"] for m in msgs if m["role"] == "system")
        user = next(m["content"] for m in msgs if m["role"] == "user")
        rows.append({"system": system, "user": user,
                     "reference": rec.get("reference", "")})
    return rows


def generate(system, user):
    payload = {"messages": [{"role": "system", "content": system},
                            {"role": "user", "content": user}], **GEN}
    r = requests.post(f"{URL}/chat/completions", json=payload, timeout=600)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"]


def run_arm(arm):
    rows = load_prompts()
    out = os.path.join(OUT_DIR, f"local_{arm}_generations.jsonl")
    n = ok = has_choices = 0
    beat_lens = []
    with open(out, "w", encoding="utf-8") as f:
        for i, row in enumerate(rows):
            gen = generate(row["system"], row["user"])
            res = validate_turn(row["system"], gen, history=row["user"])
            n += 1
            ok += int(res.ok)
            has_choices += int(bool(res.stats.get("has_choices")))
            mb = res.stats.get("max_beat_sentences")
            if mb is not None:
                beat_lens.append(mb)
            rec = {"system": row["system"], "user": row["user"],
                   "reference": row["reference"], "generated": gen,
                   "valid": res.ok, "reasons": res.reasons,
                   "max_beat_sentences": mb,
                   "has_choices": bool(res.stats.get("has_choices"))}
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            mark = "OK " if res.ok else "BAD"
            print(f"[{arm}] {i+1}/{len(rows)} {mark} "
                  f"beat={mb} choices={rec['has_choices']} "
                  f"{'; '.join(res.reasons[:1])}")
    print(f"\n[{arm}] validity {ok}/{n} ({ok/max(n,1):.0%}) · "
          f"choices {has_choices}/{n} · "
          f"mean_max_beat {sum(beat_lens)/len(beat_lens):.2f}"
          if beat_lens else f"\n[{arm}] validity {ok}/{n}")
    print(f"[{arm}] wrote {out}")


def _load(arm):
    p = os.path.join(OUT_DIR, f"local_{arm}_generations.jsonl")
    return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def compare():
    base, tuned = _load("base"), _load("tuned")
    def agg(rows):
        n = len(rows)
        ok = sum(r["valid"] for r in rows)
        ch = sum(r["has_choices"] for r in rows)
        bl = [r["max_beat_sentences"] for r in rows
              if r["max_beat_sentences"] is not None]
        return n, ok, ch, (sum(bl)/len(bl) if bl else 0)
    bn, bok, bch, bbeat = agg(base)
    tn, tok, tch, tbeat = agg(tuned)
    print("\n=== A/B: tool-call / format metrics ===")
    print(f"{'metric':<22}{'base':>10}{'tuned':>10}")
    print(f"{'format validity':<22}{f'{bok}/{bn}':>10}{f'{tok}/{tn}':>10}")
    print(f"{'has offer_choices':<22}{f'{bch}/{bn}':>10}{f'{tch}/{tn}':>10}")
    print(f"{'mean max beat (sent)':<22}{bbeat:>10.2f}{tbeat:>10.2f}")
    print("\n=== prose: per-prompt side-by-side (truncated) ===")
    for i, (b, t) in enumerate(zip(base, tuned)):
        print(f"\n--- prompt {i+1} ---")
        print(f"[BASE valid={b['valid']} beat={b['max_beat_sentences']}]"
              f"\n{b['generated'][:600]}")
        print(f"[TUNED valid={t['valid']} beat={t['max_beat_sentences']}]"
              f"\n{t['generated'][:600]}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", choices=["base", "tuned"])
    ap.add_argument("--compare", action="store_true")
    a = ap.parse_args()
    if a.compare:
        compare()
    elif a.arm:
        run_arm(a.arm)
    else:
        ap.print_help()


if __name__ == "__main__":
    raise SystemExit(main())