#!/usr/bin/env python3
"""Turn llama-bench JSON output into course lab-notebook lines.

Purpose: read the JSON llama-bench writes with `--output json`, pull out the fields that
    make a measurement meaningful (build, backend, model file, quantisation, offloaded
    layers, batch sizes, KV cache types, test sizes, repetitions, mean and standard
    deviation) and append one JSON line per test to the lab notebook, in the shape every
    other lab in this course uses. A result from the CPU backend is printed and refused, so
    a build without its GPU backend never reaches the notebook, unless --allow-cpu is given.
Platform: all (pure Python standard library; nothing platform-specific)
Minimum memory: 8 GB
Assumes: Python 3.9 or later; a JSON file produced by `llama-bench ... --output json` (the
    field names used here were read from llama-bench's JSON output at llama.cpp v0.4.0);
    the lab notebook started in Part 1, or a path to create.

Usage: python3 bench-to-labbook.py --input bench-results/run.json --labbook labbook.md
       python3 bench-to-labbook.py --input run.json --print-only

Exit status: 0 when the lines were recorded (or printed with --print-only); 1 when the input
is missing or empty, or when llama-bench reports the CPU backend and --allow-cpu was not given,
in which case nothing is written.

Every field is looked up by name and recorded as null when the build did not write it. A
null means "this build did not report it", not "it was zero". The number of repetitions
is counted from the per-repetition samples, because llama-bench's JSON has no field for it.
ctx_size is the context the test asked for, n_prompt + n_gen + n_depth, which is what
llama-bench passes to llama.cpp; llama.cpp v0.4.0 pads it up to a multiple of 256 when it
allocates (512 for pp512, 256 for tg128). ggml_cuda_enable_unified_memory records the value
of that environment variable when it was set (CUDA and HIP builds read it), and null otherwise.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from pathlib import Path

LAB = "part-06/lab-benchmark-the-reference-models"


def pick(row: dict, *names, default=None):
    """First present, non-empty value among several possible key names."""
    for name in names:
        if name in row and row[name] not in (None, ""):
            return row[name]
    return default


def as_int(value) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def test_label(row: dict) -> str:
    """llama-bench's own naming: pp512 is prompt processing, tg128 is generation."""
    n_prompt = as_int(pick(row, "n_prompt"))
    n_gen = as_int(pick(row, "n_gen"))
    depth = as_int(pick(row, "n_depth"))
    if n_prompt and n_gen:
        label = f"pp{n_prompt}+tg{n_gen}"
    elif n_gen:
        label = f"tg{n_gen}"
    elif n_prompt:
        label = f"pp{n_prompt}"
    else:
        label = "unknown"
    return f"{label}@d{depth}" if depth else label


def backend_of(row: dict) -> str | None:
    value = pick(row, "backends", "backend")
    if isinstance(value, list):
        return ",".join(str(v) for v in value)
    return str(value) if value is not None else None


# A quantisation name in a file name: Q4_K_M, Q8_0, IQ4_XS, MXFP4, BF16 ...
QUANT_RE = re.compile(r"(?:^|[-_. /])((?:IQ|TQ|Q)\d[A-Z0-9_]*|MXFP4|BF16|F16|F32)(?=$|[-_. /])")
# The same thing in llama-bench's model_type, e.g. "qwen3 8B Q4_K - Medium".
TYPE_RE = re.compile(r"\b((?:IQ|TQ|Q)\d[A-Z0-9_]*|MXFP4|BF16|F16|F32)\b(?: - (SMALL|MEDIUM|LARGE))?")
SUFFIX = {"SMALL": "_S", "MEDIUM": "_M", "LARGE": "_L"}


def quant_of(row: dict, model_path: str) -> str | None:
    """Prefer the file name, which carries the full name (Q4_K_M rather than Q4_K)."""
    for value in (model_path, pick(row, "model_filename")):
        if value:
            found = QUANT_RE.search(Path(str(value)).name.upper())
            if found:
                return found.group(1)
    model_type = pick(row, "model_type")
    if model_type:
        found = TYPE_RE.search(str(model_type).upper())
        if found:
            return found.group(1) + SUFFIX.get(found.group(2) or "", "")
    return None


def read_json(path: Path):
    """PowerShell on Windows can write UTF-16 or UTF-8 with a byte-order mark; accept both."""
    raw = path.read_bytes()
    if raw[:2] in (b"\xff\xfe", b"\xfe\xff"):
        text = raw.decode("utf-16")
    else:
        text = raw.decode("utf-8-sig")
    try:
        return json.loads(text)
    except json.JSONDecodeError as exc:
        raise SystemExit(
            f"{path}: not valid JSON ({exc}). An empty or half-written file means llama-bench "
            "failed; read the .log file beside it."
        )


def rows_from(path: Path) -> list[dict]:
    data = read_json(path)
    if isinstance(data, dict):
        for key in ("results", "data", "rows"):
            if isinstance(data.get(key), list):
                return data[key]
        return [data]
    if isinstance(data, list):
        return data
    raise SystemExit(f"{path}: expected a JSON array of test results, got {type(data).__name__}")


def record(row: dict, args) -> dict:
    size_bytes = as_int(pick(row, "model_size"))
    n_params = as_int(pick(row, "model_n_params"))
    samples = pick(row, "samples_ts", "samples_ns", default=[])
    model_path = args.model_path or pick(row, "model_filename")
    n_prompt, n_gen, n_depth = (as_int(pick(row, name)) for name in ("n_prompt", "n_gen", "n_depth"))
    return {
        "lab": LAB,
        "run": args.run_id,
        "host": args.host,
        "engine": "llama.cpp",
        "build": pick(row, "build_commit"),
        "build_number": pick(row, "build_number"),
        "backend": backend_of(row),
        "gpu_info": pick(row, "gpu_info"),
        "cpu_info": pick(row, "cpu_info"),
        "model": pick(row, "model_type", "model_filename"),
        "model_path": model_path,
        "quant": quant_of(row, args.model_path),
        "params_b": round(n_params / 1e9, 2) if n_params else None,
        "file_gb": round(size_bytes / 1e9, 2) if size_bytes else None,
        "model_size_bytes": size_bytes or None,
        "test": test_label(row),
        "n_prompt": pick(row, "n_prompt"),
        "n_gen": pick(row, "n_gen"),
        "n_depth": pick(row, "n_depth"),
        "ctx_size": n_prompt + n_gen + n_depth,
        "n_gpu_layers": pick(row, "n_gpu_layers"),
        "split_mode": pick(row, "split_mode"),
        "n_batch": pick(row, "n_batch"),
        "n_ubatch": pick(row, "n_ubatch"),
        "n_threads": pick(row, "n_threads"),
        "flash_attn": pick(row, "flash_attn"),
        "type_k": pick(row, "type_k"),
        "type_v": pick(row, "type_v"),
        "ggml_cuda_enable_unified_memory": os.environ.get("GGML_CUDA_ENABLE_UNIFIED_MEMORY"),
        "reps": len(samples) if isinstance(samples, list) and samples else None,
        "tokens_per_s": pick(row, "avg_ts"),
        "tokens_per_s_stddev": pick(row, "stddev_ts"),
        "measured_on": pick(row, "test_time"),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--input", required=True, help="JSON file written by llama-bench --output json")
    parser.add_argument("--labbook", default="labbook.md", help="notebook to append to")
    parser.add_argument("--run-id", default="", help="identifier shared by one sweep")
    parser.add_argument("--host", default="", help="short machine description")
    parser.add_argument("--model-path", default="", help="path of the model file benchmarked")
    parser.add_argument("--print-only", action="store_true", help="print the lines, record nothing")
    parser.add_argument("--allow-cpu", action="store_true",
                        help="record a result from the CPU backend (a deliberate CPU run only)")
    args = parser.parse_args()

    path = Path(args.input)
    if not path.exists():
        print(f"{path} does not exist", file=sys.stderr)
        return 1

    rows = rows_from(path)
    if not rows:
        print(f"{path} contained no test results", file=sys.stderr)
        return 1

    entries = [record(row, args) for row in rows]

    for entry in entries:
        rate, spread = entry["tokens_per_s"], entry["tokens_per_s_stddev"]
        rate_text = f"{float(rate):.2f}" if isinstance(rate, (int, float)) else "n/a"
        spread_text = f"{float(spread):.2f}" if isinstance(spread, (int, float)) else "n/a"
        print(f"    {entry['test']:>14}  {rate_text:>10} ± {spread_text:<7} tok/s"
              f"   backend {entry['backend'] or 'unknown'}")

    on_cpu = "CPU" in {entry["backend"] for entry in entries}
    if on_cpu:
        print("    WARNING: llama-bench reports the CPU backend. This build has no GPU backend, or"
              " the GPU runtime did not load (install lesson, 'Proving the backend is in use').",
              file=sys.stderr)

    if args.print_only:
        return 0

    if on_cpu and not args.allow_cpu:
        print("    nothing recorded; pass --allow-cpu only for a deliberate CPU run", file=sys.stderr)
        return 1

    notebook = Path(args.labbook)
    if not notebook.exists():
        print(f"    {notebook} does not exist; creating it", file=sys.stderr)
    with notebook.open("a", encoding="utf-8") as handle:
        for entry in entries:
            handle.write(json.dumps(entry) + "\n")
    print(f"    recorded {len(entries)} line(s) in {notebook}")
    return 0


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