#!/usr/bin/env python3
"""Turn Part 9 load-test notebook lines into curves, comparisons and a pass/fail check.

Purpose: read the JSON lines load-test.py appended to the lab notebook and print, for one
    label, the throughput and latency curve with the derived numbers the lab asks for
    (scaling efficiency, the knee, a straight-line fit of time per output token against
    concurrency, and the operating point under the recorded objective); for two labels,
    a level-by-level comparison; and for the whole notebook, the lab's validation rules.
Platform: all (spark, strix, mac, nvidia). Python 3.9 or later, standard library only.
Minimum memory: none worth stating; it reads a text file.
Assumes: a notebook (default labbook.md) holding lines written by load-test.py, whose
    "lab" field is part-09/lab-serve-twenty-concurrent-users. Other lines, prose and
    other labs' records are skipped. When a label was run more than once, the most
    recent line per concurrency level is used.

Usage:
    python3 summarise-load.py --list
    python3 summarise-load.py --label llama-baseline
    python3 summarise-load.py --label llama-baseline --markdown --record
    python3 summarise-load.py --compare llama-baseline vllm-baseline
    python3 summarise-load.py --check
"""

from __future__ import annotations

import argparse
import json
import sys
import time

LAB_ID = "part-09/lab-serve-twenty-concurrent-users"
BASELINE_LEVELS = {1, 5, 10, 20}


def load_records(path: str) -> list[dict]:
    """Every load-test row for this lab, in file order."""
    rows = []
    try:
        with open(path, encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line.startswith("{"):
                    continue
                try:
                    obj = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if obj.get("lab") == LAB_ID and "concurrency" in obj and "ttft_s" in obj:
                    rows.append(obj)
    except FileNotFoundError:
        sys.exit(f"summarise-load: {path} not found; run load-test.py first or pass --labbook")
    return rows


def by_label(rows: list[dict]) -> dict[str, dict[int, dict]]:
    """label -> concurrency -> the most recent row."""
    out: dict[str, dict[int, dict]] = {}
    for row in rows:
        out.setdefault(str(row.get("label")), {})[int(row["concurrency"])] = row
    return out


def fit_line(xs: list[float], ys: list[float]):
    """Least-squares y = a + b x. Returns (a, b) or None with fewer than two distinct x."""
    if len(set(xs)) < 2:
        return None
    n = len(xs)
    mx, my = sum(xs) / n, sum(ys) / n
    sxx = sum((x - mx) ** 2 for x in xs)
    b = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sxx
    return my - b * mx, b


def curve(levels: dict[int, dict]) -> dict:
    """Derived numbers for one label's curve."""
    cs = sorted(c for c, r in levels.items() if r.get("completed"))
    derived = {"levels": cs, "knee": None, "operating_point": None, "fit": None}
    if not cs:
        return derived
    first = levels[cs[0]]
    per_request = first["output_tokens_per_s"] / cs[0]
    derived["per_request_tokens_per_s"] = per_request
    for prev, cur in zip(cs, cs[1:]):
        gain = (levels[cur]["output_tokens_per_s"] - levels[prev]["output_tokens_per_s"]) / (cur - prev)
        if gain < 0.5 * per_request:
            derived["knee"] = [prev, cur]
            break
    for c in cs:
        r = levels[c]
        if (r["failed"] == 0 and r["ttft_s"]["p90"] <= r.get("slo_ttft_s", 2.0)
                and r["tpot_s"]["p90"] <= r.get("slo_tpot_s", 0.1)):
            derived["operating_point"] = c
    slots = first.get("server_slots")
    fit_cs = [c for c in cs if not slots or c <= slots]
    line = fit_line([float(c) for c in fit_cs], [levels[c]["tpot_s"]["p50"] for c in fit_cs])
    if line:
        a, b = line
        derived["fit"] = {"a_s": a, "b_s": b, "levels_used": fit_cs}
    return derived


def bar(value: float, top: float, width: int = 24) -> str:
    filled = int(round(width * value / top)) if top > 0 else 0
    return "#" * filled + "." * (width - filled)


def print_curve(label: str, levels: dict[int, dict], markdown: bool) -> dict:
    d = curve(levels)
    cs = d["levels"]
    if not cs:
        print(f"{label}: no level completed any request")
        return d
    first = levels[cs[0]]
    print(f"==> {label}: {first.get('engine')} {first.get('engine_version')}, model {first.get('model')}, "
          f"prompt set {first.get('prompt_set')}, max_tokens {first.get('max_tokens')}")
    print(f"    host {first.get('host', 'unknown')}, quant {first.get('quant', 'unknown')}, "
          f"context {first.get('context_length')}, server slots {first.get('server_slots')}, "
          f"SLO TTFT <= {first.get('slo_ttft_s')} s, TPOT <= {first.get('slo_tpot_s')} s")
    head = ["c", "ok", "out tok/s", "scaling", "req/s", "goodput", "TTFT p50", "TTFT p90",
            "TPOT p50", "TPOT p90", "in flight", "cached"]
    table = []
    for c in cs:
        r = levels[c]
        scaling = r["output_tokens_per_s"] / (c * d["per_request_tokens_per_s"])
        table.append([str(c), f"{r['completed']}/{r['requests']}", f"{r['output_tokens_per_s']:.1f}",
                      f"{scaling:.2f}", f"{r['requests_per_s']:.3f}",
                      f"{r.get('goodput_requests_per_s', 0.0):.3f}",
                      f"{r['ttft_s']['p50']:.3f}", f"{r['ttft_s']['p90']:.3f}",
                      f"{r['tpot_s']['p50']:.4f}", f"{r['tpot_s']['p90']:.4f}",
                      f"{r.get('mean_in_flight', 0.0):.1f}", str(r.get("cached_tokens_mean"))])
    if markdown:
        print("| " + " | ".join(head) + " |")
        print("|" + "---|" * len(head))
        for row in table:
            print("| " + " | ".join(row) + " |")
    else:
        widths = [max(len(h), *(len(row[i]) for row in table)) for i, h in enumerate(head)]
        print("    " + "  ".join(h.rjust(w) for h, w in zip(head, widths)))
        for row in table:
            print("    " + "  ".join(v.rjust(w) for v, w in zip(row, widths)))

    top_out = max(levels[c]["output_tokens_per_s"] for c in cs)
    top_tpot = max(levels[c]["tpot_s"]["p90"] for c in cs)
    print("\n    output tokens/s                     TPOT p90 (s)")
    for c in cs:
        r = levels[c]
        print(f"    c={c:>3} {bar(r['output_tokens_per_s'], top_out)} {r['output_tokens_per_s']:>8.1f}"
              f"   {bar(r['tpot_s']['p90'], top_tpot, 16)} {r['tpot_s']['p90']:.4f}")

    print()
    if d["knee"]:
        a, b = d["knee"]
        print(f"    knee: between c={a} and c={b}, each added request gained less than half "
              f"of one request's throughput at c={cs[0]}")
    else:
        print("    knee: not reached at the levels run; every step still gained at least half "
              "of one request's throughput")
    if d["fit"]:
        a_s, b_s = d["fit"]["a_s"], d["fit"]["b_s"]
        print(f"    TPOT p50 ~ {a_s * 1000:.2f} ms + {b_s * 1000:.3f} ms x c  "
              f"(fitted on c = {', '.join(map(str, d['fit']['levels_used']))})")
        for c in d["fit"]["levels_used"]:
            predicted = c / (a_s + b_s * c) if a_s + b_s * c > 0 else float("nan")
            print(f"      c={c:>3}: c / fitted TPOT = {predicted:8.1f} tok/s if every request were "
                  f"always decoding; measured {levels[c]['output_tokens_per_s']:8.1f} tok/s")
        if b_s > 0:
            slo = first.get("slo_tpot_s", 0.1)
            print(f"    ceiling as c grows: 1 / b = {1 / b_s:.1f} tok/s; "
                  f"TPOT p50 reaches {slo} s at c = {(slo - a_s) / b_s:.1f}")
        else:
            print("    b <= 0: TPOT did not rise with concurrency at these levels, so no ceiling is implied")
    if not d["fit"]:
        print("    TPOT fit: needs at least two levels at or below the server's slot count")
    print(f"    operating point (no failures, both p90s within the SLO): "
          f"{'c=' + str(d['operating_point']) if d['operating_point'] else 'none of the levels run'}")
    return d


def record_summary(path: str, label: str, levels: dict[int, dict], d: dict) -> None:
    first = levels[d["levels"][0]]
    line = {
        "lab": LAB_ID, "record": "curve-summary", "label": label,
        "engine": first.get("engine"), "engine_version": first.get("engine_version"),
        "host": first.get("host"), "quant": first.get("quant"), "model": first.get("model"),
        "context_length": first.get("context_length"), "server_slots": first.get("server_slots"),
        "levels": d["levels"], "knee_between": d["knee"], "operating_point": d["operating_point"],
        "peak_output_tokens_per_s": max(levels[c]["output_tokens_per_s"] for c in d["levels"]),
        "tpot_fit_ms": ({"a": round(d["fit"]["a_s"] * 1000, 3), "b": round(d["fit"]["b_s"] * 1000, 4)}
                        if d["fit"] else None),
        "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }
    with open(path, "a", encoding="utf-8") as handle:
        handle.write(json.dumps(line, sort_keys=True) + "\n")
    print(f"    appended a curve-summary line for {label} to {path}")


def print_compare(a: str, b: str, labels: dict[str, dict[int, dict]]) -> int:
    for name in (a, b):
        if name not in labels:
            print(f"No rows for label {name}. Labels present: {', '.join(sorted(labels)) or 'none'}",
                  file=sys.stderr)
            return 1
    shared = sorted(set(labels[a]) & set(labels[b]))
    if not shared:
        print("The two labels share no concurrency level.", file=sys.stderr)
        return 1
    ra, rb = labels[a][shared[0]], labels[b][shared[0]]
    for field in ("model", "max_tokens", "prompt_set", "context_length", "server_slots", "quant", "engine"):
        if ra.get(field) != rb.get(field):
            print(f"    note: {field} differs: {ra.get(field)} vs {rb.get(field)}")
    print(f"==> {a} (A) against {b} (B)")
    head = ["c", "out/s A", "out/s B", "B/A", "goodput A", "goodput B", "TTFT p50 A", "TTFT p50 B",
            "TTFT p90 A", "TTFT p90 B", "TPOT p50 A", "TPOT p50 B", "cached A", "cached B"]
    rows = []
    for c in shared:
        x, y = labels[a][c], labels[b][c]
        ratio = y["output_tokens_per_s"] / x["output_tokens_per_s"] if x["output_tokens_per_s"] else float("nan")
        rows.append([str(c), f"{x['output_tokens_per_s']:.1f}", f"{y['output_tokens_per_s']:.1f}", f"{ratio:.2f}",
                     f"{x.get('goodput_requests_per_s', 0.0):.3f}", f"{y.get('goodput_requests_per_s', 0.0):.3f}",
                     f"{x['ttft_s']['p50']:.3f}", f"{y['ttft_s']['p50']:.3f}",
                     f"{x['ttft_s']['p90']:.3f}", f"{y['ttft_s']['p90']:.3f}",
                     f"{x['tpot_s']['p50']:.4f}", f"{y['tpot_s']['p50']:.4f}",
                     str(x.get("cached_tokens_mean")), str(y.get("cached_tokens_mean"))])
    widths = [max(len(h), *(len(r[i]) for r in rows)) for i, h in enumerate(head)]
    print("    " + "  ".join(h.rjust(w) for h, w in zip(head, widths)))
    for r in rows:
        print("    " + "  ".join(v.rjust(w) for v, w in zip(r, widths)))
    return 0


def check(labels: dict[str, dict[int, dict]]) -> int:
    results = []

    def verdict(ok: bool, text: str, level: str = "FAIL") -> None:
        results.append("PASS" if ok else level)
        print(f"    {'PASS' if ok else level:<4}  {text}")

    print("==> checking the notebook against the lab's validation rules")
    sweeps = [name for name, lv in labels.items()
              if BASELINE_LEVELS <= set(lv) and all(lv[c]["completed"] for c in BASELINE_LEVELS)
              and lv[1].get("prompt_set") == "mixed"]
    verdict(bool(sweeps), f"a mixed-prompt sweep with every level of {sorted(BASELINE_LEVELS)} completed: "
            f"{', '.join(sweeps) or 'none'}")
    for name in sweeps:
        lv = labels[name]
        failed = sum(lv[c]["failed"] for c in lv)
        verdict(failed == 0, f"{name}: failed requests across all levels = {failed}", "WARN")
        verdict(all(lv[c].get("token_counts") == "usage" for c in lv),
                f"{name}: output tokens counted from the server's usage block", "WARN")
        missing = [k for k in ("host", "quant", "context_length", "server_slots")
                   if lv[1].get(k) in (None, "unknown")]
        verdict(not missing, f"{name}: context recorded ({', '.join(missing) + ' missing' if missing else 'complete'})",
                "WARN")
    shared = [n for n, lv in labels.items() if any(r.get("prompt_set") == "shared-prefix" for r in lv.values())]
    unique = [n for n, lv in labels.items() if any(r.get("prompt_set") == "unique-prefix" for r in lv.values())]
    verdict(bool(shared) and bool(unique), f"prefix pair present: shared {shared or 'none'}, unique {unique or 'none'}")
    for name in shared:
        for c, r in sorted(labels[name].items()):
            if r.get("cached_tokens_mean") is None or not r.get("prompt_tokens_mean"):
                verdict(False, f"{name} c={c}: the server reported no cached-token count", "WARN")
                continue
            frac = r["cached_tokens_mean"] / r["prompt_tokens_mean"]
            verdict(frac > 0.5, f"{name} c={c}: mean cached tokens are {frac:.0%} of the prompt")
    for name in unique:
        for c, r in sorted(labels[name].items()):
            if r.get("cached_tokens_mean") is not None and r.get("prompt_tokens_mean"):
                frac = r["cached_tokens_mean"] / r["prompt_tokens_mean"]
                verdict(frac < 0.1, f"{name} c={c}: mean cached tokens are {frac:.0%} of the prompt (control)")
    fails = results.count("FAIL")
    print(f"    {results.count('PASS')} pass, {results.count('WARN')} warn, {fails} fail")
    return 1 if fails else 0


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--list", action="store_true", help="list the labels in the notebook")
    group.add_argument("--label", help="print the curve for one label")
    group.add_argument("--compare", nargs=2, metavar=("A", "B"), help="compare two labels level by level")
    group.add_argument("--check", action="store_true", help="run the lab's validation rules")
    parser.add_argument("--markdown", action="store_true", help="print the curve table as Markdown")
    parser.add_argument("--record", action="store_true",
                        help="with --label, append a curve-summary line to the notebook")
    args = parser.parse_args(argv)

    labels = by_label(load_records(args.labbook))
    if args.list:
        print(f"==> {len(labels)} label(s) in {args.labbook}")
        for name, lv in labels.items():
            r = lv[max(lv)]
            print(f"    {name:<28} {r.get('engine', '?'):<10} {r.get('prompt_set', '?'):<14} "
                  f"levels {','.join(map(str, sorted(lv)))}  last {r.get('recorded_at', '?')}")
        return 0
    if args.check:
        return check(labels)
    if args.compare:
        return print_compare(args.compare[0], args.compare[1], labels)
    if args.label not in labels:
        print(f"No rows for label {args.label}. Labels present: {', '.join(sorted(labels)) or 'none'}",
              file=sys.stderr)
        return 1
    d = print_curve(args.label, labels[args.label], args.markdown)
    if args.record and d["levels"]:
        record_summary(args.labbook, args.label, labels[args.label], d)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))