#!/usr/bin/env python3
"""Turn this lab's notebook lines into the course benchmark tables, and check them.

Purpose: read the JSON lines that compare-engines.py and feature-probe.py appended to the lab
    notebook, keep the most recent measurement per engine, and print (1) a Markdown table of
    speed and memory, (2) a Markdown table of API features, and (3) the rows in the form the
    course's <Benchmark> recording sheet takes. With --check, test every line against the
    lab's validation rules and exit non-zero if any rule fails.
Platform: all (Python standard library only)
Minimum memory: 12 GB
Assumes: Python 3.9 or later; a notebook file in which each result is one JSON object per line
    (other lines, such as Markdown notes, are ignored).

Usage: python3 summarise-engines.py --labbook labbook.md
       python3 summarise-engines.py --labbook labbook.md --check
       python3 summarise-engines.py --labbook labbook.md --check --allow-cpu   (script testing only)
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

LAB = "part-08/lab-same-model-every-engine"
CACHE_LIMIT = 64  # more reused prompt tokens than this means the run measured the prefix cache


def load(path: Path) -> list:
    records = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if record.get("lab") == LAB:
            records.append(record)
    return records


def latest(records: list) -> tuple:
    """Most recent short and long speed line, and most recent feature line, per engine."""
    speed, features, order = {}, {}, []
    for r in sorted(records, key=lambda x: x.get("measured_on", "")):
        engine = r.get("engine", "?")
        if engine not in order:
            order.append(engine)
        if r.get("probe") == "features":
            features[engine] = r
        elif r.get("test") in ("short", "long"):
            speed.setdefault(engine, {})[r["test"]] = r
    return order, speed, features


def memory_added(record: dict) -> tuple:
    """The memory reading that best describes what this server added, and its name.

    A discrete NVIDIA card reports a device total, and nvidia-smi's per-process figure counts
    every CUDA allocation of the server wherever it runs. Under WSL2 NVML does not report
    compute processes (CUDA on WSL guide, features not yet supported), so a per-process delta
    that is missing, zero or negative falls back to the device figure. An AMD GPU reports VRAM
    and GTT use, which together cover the BIOS carve-out and the dynamically mapped share. A
    DGX Spark reports no device total, and NVIDIA's guidance there is MemAvailable. A Mac has
    vm_stat.
    """
    delta = record.get("memory_delta_mb") or {}
    loaded = record.get("memory") or {}
    if loaded.get("gpu_total_mb"):
        if delta.get("gpu_process_mb") is not None and delta["gpu_process_mb"] > 0:
            return delta["gpu_process_mb"], "nvidia-smi per-process"
        if delta.get("gpu_used_mb") is not None:
            return delta["gpu_used_mb"], "nvidia-smi device"
    if "rocm_vram_used_mb" in delta or "rocm_gtt_used_mb" in delta:
        return delta.get("rocm_vram_used_mb", 0) + delta.get("rocm_gtt_used_mb", 0), "rocm-smi VRAM+GTT"
    if delta.get("mem_available_mb") is not None:
        return delta["mem_available_mb"], "MemAvailable"
    if delta.get("mac_used_mb") is not None:
        return delta["mac_used_mb"], "vm_stat active+wired+compressed"
    return None, "no idle snapshot"


def fmt(value, digits=2) -> str:
    if value is None:
        return "—"
    if isinstance(value, float):
        return f"{value:.{digits}f}"
    return str(value)


def yes_no(section: dict) -> str:
    if not section:
        return "—"
    return "yes" if section.get("supported") else "no"


def print_tables(order, speed, features) -> None:
    print("Speed and memory (medians; prefill from the difference between the two prompts)\n")
    print("| Engine | Quant | Backend | TTFT short s | TTFT long s | Prefill tok/s | "
          "Decode short tok/s | Decode long tok/s | Memory added MB | Memory source |")
    print("| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |")
    rows = []
    for engine in order:
        pair = speed.get(engine)
        if not pair:
            continue
        s, l = pair.get("short", {}), pair.get("long", {})
        any_line = l or s
        mem, source = memory_added(any_line)
        print(f"| {engine} | {any_line.get('quant') or '—'} | {any_line.get('backend') or '—'} | "
              f"{fmt(s.get('ttft_s'), 3)} | {fmt(l.get('ttft_s'), 3)} | "
              f"{fmt(any_line.get('prefill_tokens_per_s'), 1)} | {fmt(s.get('decode_tokens_per_s'))} | "
              f"{fmt(l.get('decode_tokens_per_s'))} | {fmt(mem)} | {source} |")
        rows.append([f"{engine}, {any_line.get('quant') or '?'}", s.get("ttft_s"), l.get("ttft_s"),
                     any_line.get("prefill_tokens_per_s"), s.get("decode_tokens_per_s"),
                     l.get("decode_tokens_per_s"), mem])

    print("\nAPI features\n")
    print("| Engine | /v1/models | Tool calling | JSON schema | Detail |")
    print("| --- | --- | --- | --- | --- |")
    feature_rows = []
    for engine in order:
        f = features.get(engine)
        if not f:
            continue
        details = [d for d in (f.get("tool_calling", {}).get("detail"),
                               f.get("structured_output", {}).get("detail")) if d]
        detail = "; ".join(details) or "—"
        cells = [yes_no(f.get("models_endpoint")), yes_no(f.get("tool_calling")),
                 yes_no(f.get("structured_output"))]
        print(f"| {engine} | {cells[0]} | {cells[1]} | {cells[2]} | {detail} |")
        feature_rows.append([engine, *cells, detail])

    def js(value):
        return "'—'" if value is None else (json.dumps(value) if not isinstance(value, str) else repr(value))

    print("\nRows for the speed recording sheet (paste into the Benchmark rows prop)\n")
    for row in rows:
        print("  [" + ", ".join(js(v) for v in row) + "],")
    print("\nRows for the features recording sheet\n")
    for row in feature_rows:
        print("  [" + ", ".join(js(v) for v in row) + "],")


def check(order, speed, features, allow_cpu: bool) -> int:
    failures = 0

    def report(level: str, engine: str, message: str) -> None:
        nonlocal failures
        failures += level == "FAIL"
        print(f"{level:<4}  {engine:<16} {message}")

    if not order:
        report("FAIL", "-", f"no lines with lab = {LAB}")
    for engine in order:
        pair = speed.get(engine, {})
        if set(pair) != {"short", "long"}:
            report("FAIL", engine, f"speed tests recorded: {sorted(pair) or 'none'} (need short and long)")
        for test, r in sorted(pair.items()):
            label = f"{test}:"
            if r.get("ttft_s") is None or r.get("decode_tokens_per_s") is None:
                report("FAIL", engine, f"{label} ttft_s or decode_tokens_per_s is null")
            if not r.get("engine_version") or not r.get("quant"):
                report("FAIL", engine, f"{label} engine_version or quant is empty")
            backend = (r.get("backend") or "").strip()
            if not backend or (backend.upper() == "CPU" and not allow_cpu):
                report("FAIL", engine, f"{label} backend is '{backend}'; it must name the accelerator")
            if r.get("prompt_tag") is False:
                report("FAIL", engine, f"{label} recorded with --reuse-prompt; that measures the cache")
            cached = r.get("max_cached_prompt_tokens")
            if cached is None:
                report("WARN", engine, f"{label} server reported no cached-token count; check its log")
            elif cached > CACHE_LIMIT:
                report("FAIL", engine, f"{label} {cached} prompt tokens came from the cache")
            if r.get("completion_tokens") != r.get("max_tokens"):
                report("WARN", engine, f"{label} {r.get('completion_tokens')} tokens generated, "
                                       f"not {r.get('max_tokens')} (finish_reason {r.get('finish_reason')})")
            if r.get("token_source") != "usage":
                report("WARN", engine, f"{label} token counts estimated from chunks")
        any_line = pair.get("long") or pair.get("short")
        if any_line:
            if any_line.get("prefill_tokens_per_s") is None:
                report("WARN", engine, "prefill from the difference could not be computed")
            if not any_line.get("memory_delta_mb"):
                report("WARN", engine, "no memory delta (was --idle-snapshot given?)")
            else:
                mem, source = memory_added(any_line)
                if mem is not None and mem <= 0:
                    report("WARN", engine, f"memory added is {mem} MB ({source}); a loaded model "
                                           "cannot add nothing, so take a new idle snapshot")
        f = features.get(engine)
        if not f:
            report("FAIL", engine, "no feature-probe line")
        else:
            for section in ("tool_calling", "structured_output"):
                if not isinstance((f.get(section) or {}).get("supported"), bool):
                    report("FAIL", engine, f"feature line has no boolean {section}.supported")
        if pair and f and set(pair) == {"short", "long"}:
            report("PASS", engine, "speed pair and feature line present")
    print(f"\n{failures} failure(s)")
    return 1 if failures else 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md", help="notebook to read")
    parser.add_argument("--check", action="store_true", help="validate the lines instead of printing tables")
    parser.add_argument("--allow-cpu", action="store_true", help="accept a CPU backend (for testing the scripts)")
    args = parser.parse_args()

    path = Path(args.labbook)
    if not path.exists():
        print(f"{path} does not exist", file=sys.stderr)
        return 1
    order, speed, features = latest(load(path))
    if args.check:
        return check(order, speed, features, args.allow_cpu)
    print_tables(order, speed, features)
    return 0


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