#!/usr/bin/env python3
"""Measure one generation on a split cluster: prefill, decode, and bytes over the link.

Purpose: send one completion to a llama-server that is running across RPC hosts, read the
    prefill and decode timings the server reports, read this machine's interface byte
    counters either side of the request, and append one lab-notebook line describing the
    run. The point is not the tokens-per-second figure on its own: it is that figure next
    to the number of bytes the cluster link carried to produce it.
Platform: all (Linux reads /proc/net/dev, macOS reads netstat -ib; nothing else is needed)
Minimum memory: 8 GB on the client; the model has to fit the sum of the cluster's devices
Assumes: Python 3.9 or later; a llama-server started by run-split.sh and still loading or
    loaded; the interface name from CLUSTER_IFACE; the lab notebook from Part 1.

Usage: python3 measure-split.py --iface enp1s0f1np1 --labbook labbook.md
       python3 measure-split.py --iface en5 --server-url http://127.0.0.1:8080 \
           --prompt-words 400 --n-predict 128 --note "two Sparks over RoCE" \
           --transport rdma --tensor-split 1,1

Counters are for the whole interface, so anything else using that link during the run is
counted too. Run it on a quiet cluster, and treat the byte figures as the order of
magnitude they are, not as an exact accounting of the model's traffic.
"""

from __future__ import annotations

import argparse
import json
import platform
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

PROMPT_SEED = (
    "A cluster is two machines and a cable, and the cable is the part that decides. "
    "Explain, carefully and at length, what a layer split sends between machines and why. "
)


def counters(iface: str) -> tuple[int | None, int | None]:
    """Bytes received and transmitted on one interface, or (None, None) if unreadable."""
    system = platform.system()
    if system == "Linux":
        try:
            for line in Path("/proc/net/dev").read_text(encoding="utf-8").splitlines():
                name, _, rest = line.partition(":")
                if name.strip() != iface:
                    continue
                fields = rest.split()
                return int(fields[0]), int(fields[8])
        except (OSError, ValueError, IndexError):
            return None, None
        return None, None
    if system == "Darwin":
        try:
            out = subprocess.run(
                ["netstat", "-ib"], capture_output=True, text=True, check=True, timeout=20
            ).stdout
        except (OSError, subprocess.SubprocessError):
            return None, None
        for line in out.splitlines():
            fields = line.split()
            # Name Mtu Network Address Ipkts Ierrs Ibytes Opkts Oerrs Obytes Coll
            if len(fields) >= 10 and fields[0] == iface:
                try:
                    return int(fields[6]), int(fields[9])
                except ValueError:
                    continue
        return None, None
    return None, None


def get_json(url: str, timeout: float = 10.0):
    try:
        with urllib.request.urlopen(url, timeout=timeout) as handle:  # noqa: S310 - local server
            return json.loads(handle.read().decode("utf-8"))
    except (urllib.error.URLError, OSError, ValueError):
        return None


def get_text(url: str, timeout: float = 10.0) -> str | None:
    try:
        with urllib.request.urlopen(url, timeout=timeout) as handle:  # noqa: S310 - local server
            return handle.read().decode("utf-8")
    except (urllib.error.URLError, OSError):
        return None


def post_completion(base: str, prompt: str, n_predict: int, timeout: float):
    body = json.dumps(
        {"prompt": prompt, "n_predict": n_predict, "temperature": 0.0, "cache_prompt": False}
    ).encode("utf-8")
    request = urllib.request.Request(
        f"{base}/completion", data=body, headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(request, timeout=timeout) as handle:  # noqa: S310 - local server
        return json.loads(handle.read().decode("utf-8"))


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


def metrics_of(text: str | None) -> dict:
    """Pull the llamacpp: gauges out of the Prometheus exposition text."""
    wanted = {
        "llamacpp:prompt_tokens_seconds": "metrics_prompt_tokens_per_s",
        "llamacpp:predicted_tokens_seconds": "metrics_predicted_tokens_per_s",
        "llamacpp:n_decode_total": "metrics_decode_calls_total",
    }
    found: dict = {}
    if not text:
        return found
    for line in text.splitlines():
        if line.startswith("#") or " " not in line:
            continue
        name, _, value = line.partition(" ")
        key = wanted.get(name.split("{")[0])
        if key:
            try:
                found[key] = float(value)
            except ValueError:
                pass
    return found


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--server-url", default="http://127.0.0.1:8080", help="llama-server base URL")
    parser.add_argument("--iface", default="", help="cluster interface to read counters from")
    parser.add_argument("--prompt-words", type=int, default=400, help="approximate prompt length")
    parser.add_argument("--n-predict", type=int, default=128, help="tokens to generate")
    parser.add_argument("--timeout", type=float, default=900.0, help="seconds to wait for the reply")
    parser.add_argument("--labbook", default="labbook.md", help="notebook to append to")
    parser.add_argument("--lab", default="part-19/lab-a-model-bigger-than-any-one-machine")
    parser.add_argument("--note", default="", help="free text: topology, transport, anything")
    parser.add_argument("--transport", default="", help="tcp, rdma or unknown")
    parser.add_argument("--tensor-split", default="", help="the proportions this run used")
    parser.add_argument("--print-only", action="store_true", help="print it, record nothing")
    args = parser.parse_args()

    base = args.server_url.rstrip("/")
    props = get_json(f"{base}/props")
    if props is None:
        print(f"no llama-server answering at {base}; start run-split.sh first", file=sys.stderr)
        return 1

    words = PROMPT_SEED.split()
    prompt = " ".join((words * (args.prompt_words // len(words) + 1))[: args.prompt_words])

    rx0, tx0 = counters(args.iface) if args.iface else (None, None)
    wall0 = time.monotonic()
    try:
        reply = post_completion(base, prompt, args.n_predict, args.timeout)
    except (urllib.error.URLError, OSError, ValueError) as exc:
        print(f"the completion request failed: {exc}", file=sys.stderr)
        print("On a cluster this usually means a host went away mid-generation.", file=sys.stderr)
        return 1
    wall = time.monotonic() - wall0
    rx1, tx1 = counters(args.iface) if args.iface else (None, None)

    timings = reply.get("timings") or {}
    entry = {
        "record": "split-run",
        "lab": args.lab,
        "engine": "llama.cpp",
        "build": pick(props.get("build_info", {}) if isinstance(props.get("build_info"), dict) else {},
                      "build", "commit") or pick(props, "build_info"),
        "model": pick(props, "model_path", "model") or reply.get("model"),
        "client": f"{platform.system()}-{platform.machine()}",
        "transport": args.transport or "unknown",
        "tensor_split": args.tensor_split or "by free memory",
        "iface": args.iface or None,
        "prompt_tokens": pick(timings, "prompt_n", "tokens_evaluated"),
        "prompt_ms": pick(timings, "prompt_ms"),
        "prompt_tokens_per_s": pick(timings, "prompt_per_second"),
        "predicted_tokens": pick(timings, "predicted_n", "tokens_predicted"),
        "predicted_ms": pick(timings, "predicted_ms"),
        "predicted_tokens_per_s": pick(timings, "predicted_per_second"),
        "wall_s": round(wall, 3),
        "rx_bytes": (rx1 - rx0) if (rx0 is not None and rx1 is not None) else None,
        "tx_bytes": (tx1 - tx0) if (tx0 is not None and tx1 is not None) else None,
        "note": args.note,
        "measured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    entry.update(metrics_of(get_text(f"{base}/metrics")))

    def show(label: str, value, unit: str = "") -> None:
        text = f"{value:,.2f}" if isinstance(value, float) else ("n/a" if value is None else f"{value:,}")
        print(f"    {label:<28} {text:>16} {unit}")

    print("==> one completion across the cluster")
    show("prompt tokens", entry["prompt_tokens"])
    show("prefill", entry["prompt_tokens_per_s"], "tokens/s")
    show("generated tokens", entry["predicted_tokens"])
    show("decode", entry["predicted_tokens_per_s"], "tokens/s")
    show("wall clock", entry["wall_s"], "s")
    show("bytes received on link", entry["rx_bytes"])
    show("bytes sent on link", entry["tx_bytes"])
    if entry["tx_bytes"] and entry["predicted_tokens"]:
        per_token = (entry["tx_bytes"] + (entry["rx_bytes"] or 0)) / float(entry["predicted_tokens"])
        show("link bytes per token", round(per_token, 1))
        print("    Compare that against the hidden-state arithmetic in this part's first lesson.")
    if entry["rx_bytes"] is None:
        print("    No counters: pass --iface with the interface name from CLUSTER_IFACE.")

    if args.print_only:
        return 0

    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:
        handle.write(json.dumps(entry) + "\n")
    print(f"    recorded 1 line in {notebook}")
    return 0


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