#!/usr/bin/env python3
"""Compare a disaggregated run against a monolithic run, from the lab notebook.

Purpose: read the JSON lines Part 9's load generator wrote, pair the disaggregated run
    with the single-machine baseline at each concurrency level, and print what the split
    did to time to first token, time per output token and throughput. Where run-load.sh
    also recorded interface byte counters, report the bytes moved per request beside the
    key-value cache size the model's shape predicts, so that "the transfer happened" is a
    number rather than a hope. Appends one comparison line to the notebook.
Platform: all. Pure Python standard library: no pip install, and it may run anywhere the
    notebook file is, including a machine that took no part in the serving.
Minimum memory: 1 GB. It reads a text file.
Assumes: Python 3.9 or later; a labbook.md containing at least one line from each label,
    written by Part 9's load-test.py through run-load.sh. The two runs must have used the
    same model, the same context length, the same prompt set and the same concurrency
    levels, or the comparison is between two different experiments.

Usage:
    python3 compare-disagg.py --labbook labbook.md \\
        --disagg disagg --baseline baseline

    python3 compare-disagg.py --labbook labbook.md \\
        --disagg disagg-nixl --baseline baseline \\
        --kv-bytes-per-token 147456 --prompt-tokens 8192 \\
        --note "two Sparks over the QSFP cable, NixlConnector"

The --kv-bytes-per-token value comes from the course model reference: it is the model's
layers x key-value heads x head dimension x 2 x 2. Qwen3-8B is 147456; Qwen3-1.7B is
114688. Passing it turns the byte counters into a prediction you can check.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path


def read_lines(path: Path) -> list[dict]:
    """Every JSON object in the notebook, in file order. Prose lines are skipped."""
    records = []
    try:
        text = path.read_text(encoding="utf-8")
    except OSError as exc:
        print(f"Could not read {path}: {exc}", file=sys.stderr)
        return records
    for raw in text.splitlines():
        raw = raw.strip()
        if not raw.startswith("{"):
            continue
        try:
            obj = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if isinstance(obj, dict):
            records.append(obj)
    return records


def load_rows(records: list[dict], label: str) -> dict[int, dict]:
    """The most recent load-generator row per concurrency level, for one label."""
    rows: dict[int, dict] = {}
    for obj in records:
        if obj.get("label") != label:
            continue
        if obj.get("record") == "link":
            continue
        if "concurrency" not in obj or "ttft_s" not in obj:
            continue
        try:
            level = int(obj["concurrency"])
        except (TypeError, ValueError):
            continue
        rows[level] = obj
    return rows


def link_line(records: list[dict], label: str) -> dict | None:
    """The most recent interface-counter line for one label, if run-load.sh wrote one."""
    found = None
    for obj in records:
        if obj.get("record") == "link" and obj.get("label") == label:
            found = obj
    return found


def ratio(new: float, old: float) -> str:
    """A readable multiple, guarding against a zero denominator."""
    if not old:
        return "n/a"
    return f"{new / old:.2f}x"


def signed_pct(new: float, old: float) -> str:
    if not old:
        return "n/a"
    return f"{(new - old) / old * 100:+.1f}%"


def human_bytes(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if abs(n) < 1000 or unit == "TB":
            return f"{n:.2f} {unit}" if unit != "B" else f"{n:.0f} B"
        n /= 1000
    return f"{n:.2f} TB"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md",
                        help="notebook file to read and append to")
    parser.add_argument("--disagg", default="disagg",
                        help="label of the disaggregated run")
    parser.add_argument("--baseline", default="baseline",
                        help="label of the single-machine run it is compared against")
    parser.add_argument("--kv-bytes-per-token", type=int, default=0,
                        help="bytes of key-value cache per token, from the model reference")
    parser.add_argument("--prompt-tokens", type=int, default=0,
                        help="approximate prompt length used, for the predicted payload")
    parser.add_argument("--note", default="",
                        help="free text: topology, connector, link, anything")
    parser.add_argument("--lab", default="part-22/lab-two-machine-prefill-decode")
    parser.add_argument("--print-only", action="store_true",
                        help="print the comparison, write nothing")
    args = parser.parse_args()

    book = Path(args.labbook)
    records = read_lines(book)
    if not records:
        print(f"No JSON lines found in {book}. Run the load generator first.", file=sys.stderr)
        return 2

    disagg = load_rows(records, args.disagg)
    base = load_rows(records, args.baseline)

    missing = []
    if not disagg:
        missing.append(args.disagg)
    if not base:
        missing.append(args.baseline)
    if missing:
        print(f"No load-generator rows for label(s): {', '.join(missing)}.", file=sys.stderr)
        print("Labels present: " + ", ".join(sorted({
            str(r.get("label")) for r in records if r.get("label")
        })), file=sys.stderr)
        return 2

    shared = sorted(set(disagg) & set(base))
    if not shared:
        print("The two runs share no concurrency level, so nothing is comparable.",
              file=sys.stderr)
        print(f"  {args.disagg}: {sorted(disagg)}", file=sys.stderr)
        print(f"  {args.baseline}: {sorted(base)}", file=sys.stderr)
        return 2

    # --- a sanity check the reader will otherwise skip ---------------------------------
    warnings = []
    for level in shared:
        d, b = disagg[level], base[level]
        for field in ("model", "prompt_set", "max_tokens"):
            if d.get(field) != b.get(field):
                warnings.append(
                    f"  concurrency {level}: {field} differs "
                    f"({d.get(field)!r} against {b.get(field)!r})"
                )
    if warnings:
        print("These two runs are not the same experiment:")
        for line in warnings:
            print(line)
        print("  Fix the settings and run both again; the comparison below is not valid.\n")

    print(f"==> {args.disagg} against {args.baseline}, from {book}")
    print(f"    {'conc':>4}  {'TTFT p50':>18}  {'TPOT p50':>18}  {'out tok/s':>18}")
    comparison = []
    for level in shared:
        d, b = disagg[level], base[level]
        d_ttft = float(d["ttft_s"]["p50"])
        b_ttft = float(b["ttft_s"]["p50"])
        d_tpot = float(d["tpot_s"]["p50"])
        b_tpot = float(b["tpot_s"]["p50"])
        d_thru = float(d["output_tokens_per_s"])
        b_thru = float(b["output_tokens_per_s"])
        print(
            f"    {level:>4}  "
            f"{ratio(d_ttft, b_ttft):>8} {signed_pct(d_ttft, b_ttft):>9}  "
            f"{ratio(d_tpot, b_tpot):>8} {signed_pct(d_tpot, b_tpot):>9}  "
            f"{ratio(d_thru, b_thru):>8} {signed_pct(d_thru, b_thru):>9}"
        )
        comparison.append({
            "concurrency": level,
            "ttft_p50_disagg_s": round(d_ttft, 4),
            "ttft_p50_baseline_s": round(b_ttft, 4),
            "tpot_p50_disagg_s": round(d_tpot, 5),
            "tpot_p50_baseline_s": round(b_tpot, 5),
            "output_tokens_per_s_disagg": round(d_thru, 2),
            "output_tokens_per_s_baseline": round(b_thru, 2),
        })

    print("\n    A ratio below 1.00 on time to first token is the split helping. A ratio")
    print("    above 1.00 is the transfer costing more than the prefill it replaced, which")
    print("    is the expected result on ordinary Ethernet and is worth recording as such.")

    # --- the link, if the counters were read ------------------------------------------
    link = link_line(records, args.disagg)
    predicted = None
    observed = None
    if args.kv_bytes_per_token and args.prompt_tokens:
        predicted = args.kv_bytes_per_token * args.prompt_tokens
    if link and link.get("rx_bytes") is not None:
        total_requests = 0
        for level in shared:
            total_requests += int(disagg[level].get("requests", 0))
        moved = max(int(link["rx_bytes"]), int(link["tx_bytes"]))
        if total_requests:
            observed = moved / total_requests

    if observed is not None:
        print(f"\n==> Link on {link.get('iface')} during the {args.disagg} run")
        print(f"    bytes moved per request, observed   {human_bytes(observed)}")
        if predicted is not None:
            print(f"    one request's cache, from arithmetic {human_bytes(predicted)}")
            print("    Same order of magnitude means the cache is crossing the link.")
            print("    Three orders smaller means it is not, and the decode instance is")
            print("    quietly prefilling every prompt itself.")
    elif link:
        print(f"\n==> The {args.disagg} run recorded no byte counters. Set CLUSTER_IFACE.")
    else:
        print(f"\n==> No link line for label {args.disagg}. Run it through run-load.sh.")

    record = {
        "lab": args.lab,
        "record": "comparison",
        "disagg_label": args.disagg,
        "baseline_label": args.baseline,
        "model": disagg[shared[0]].get("model"),
        "prompt_set": disagg[shared[0]].get("prompt_set"),
        "connector": (link or {}).get("connector"),
        "iface": (link or {}).get("iface"),
        "kv_bytes_per_token": args.kv_bytes_per_token or None,
        "prompt_tokens": args.prompt_tokens or None,
        "predicted_bytes_per_request": predicted,
        "observed_bytes_per_request": round(observed, 1) if observed is not None else None,
        "levels": comparison,
        "same_experiment": not warnings,
        "note": args.note,
        "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    }

    if args.print_only:
        print("\n    --print-only: nothing written.")
        return 0

    with book.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")
    print(f"\n    Appended one comparison line to {book}")
    return 0


if __name__ == "__main__":
    sys.exit(main())