"""Measure prefill, decode and link use on a served cluster, and record all three.

Purpose: drive an OpenAI-compatible endpoint that is being served by a cluster, and
    record the numbers that decide whether the cluster was worth building. With
    --concurrency 1 it answers the single-user question; with a larger value it
    answers the batch question, and the two are different questions.
    Prefill is measured as prompt tokens divided by time to first token on a long
    prompt. Decode is measured as generated tokens divided by the time between the
    first and last token. Link use is read from the interface byte counters before
    and after the run, so you can see how much traffic the split actually put on the
    cable rather than assuming. One JSON line per run goes to the lab notebook.
Platform: mac (Track M) for the link counters, which come from BSD `netstat -ibn`;
    the timing half runs anywhere python3 does, including against a server on another
    machine. On Linux the link section is recorded as unavailable rather than faked.
Minimum memory: none on the machine running this script. It streams text and does
    no inference of its own, so it can be run from a laptop against the cluster.
Assumes: a server already answering POST {base-url}/chat/completions with streaming,
    reachable over plain HTTP; this is a house-network tool and does not speak TLS.
    An API key, if the server wants one, is read from the environment variable named
    by --api-key-env and is never written to this file or to the notebook.

Counting note: one streamed delta carrying content counts as one output token, which
    is exact for mlx_lm.server and for exo as this course reads them. Prefill speed
    from time to first token includes queueing and template rendering, so it is a
    lower bound on the engine's own prefill rate; --exo-bench-url asks exo for its
    own figures instead, which are measured inside the server.

Usage:
    python3 measure-pair.py --label two-macs-jaccl-tensor

    python3 measure-pair.py --base-url http://127.0.0.1:8080/v1 \
        --model mlx-community/Qwen3-32B-8bit --requests 20 --max-tokens 128 \
        --iface en2 --label two-macs-ring-tensor --labbook labbook.md

    python3 measure-pair.py --base-url http://localhost:52415/v1 \
        --exo-bench-url http://localhost:52415 --label exo-jaccl

    python3 measure-pair.py --concurrency 8 --requests 40 --label batch-of-eight
"""

from __future__ import annotations

import argparse
import json
import os
import platform
import shutil
import statistics
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
from datetime import date, datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent

# Filler with no interesting answer in it, repeated to reach the prompt length the
# prefill measurement wants. The model is being timed, not questioned.
FILLER = (
    "A cluster is two machines and a cable. The weights are divided between them, "
    "the activations cross the cable, and every design question in this part is "
    "about how often that happens and how many bytes go each time. "
)

QUESTION = "In one sentence, what limits decode speed on a single machine?"

# Roughly four characters per token for English prose. This is only used to build a
# prompt of about the right length; the exact prompt token count comes back from the
# server in the usage field where the server reports one.
CHARS_PER_TOKEN = 4


def parse_args():
    p = argparse.ArgumentParser(description="Measure a served cluster")
    p.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    p.add_argument(
        "--model",
        default=None,
        help="Model name the server reports at /models. Default: ask the server.",
    )
    p.add_argument("--requests", type=int, default=20)
    p.add_argument(
        "--concurrency",
        type=int,
        default=1,
        help="Requests in flight at once. 1 is the single-user case; a larger "
        "number is the batch case, and the two answer different questions.",
    )
    p.add_argument("--max-tokens", type=int, default=128)
    p.add_argument(
        "--prefill-tokens",
        type=int,
        default=2048,
        help="Approximate prompt length for the prefill measurement.",
    )
    p.add_argument(
        "--iface",
        default=os.environ.get("LINK_IFACE", ""),
        help="Interface to read byte counters for. Empty records every interface.",
    )
    p.add_argument(
        "--exo-bench-url",
        default=None,
        help="exo base URL; use its /bench/chat/completions for server-side figures.",
    )
    p.add_argument("--api-key-env", default="LLM_API_KEY")
    p.add_argument("--label", default="", help="A name for this run, for the record.")
    p.add_argument("--labbook", default=str(HERE / "labbook.md"))
    p.add_argument("--timeout", type=float, default=600.0)
    p.add_argument(
        "--print",
        dest="print_only",
        action="store_true",
        help="Show the record and write nothing.",
    )
    return p.parse_args()


# ------------------------------------------------------------------ link counters

def read_link_counters(iface: str):
    """Bytes in and out per interface, from BSD netstat. None where unavailable."""
    if shutil.which("netstat") is None:
        return None
    cmd = ["netstat", "-ibn"]
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=20, check=False)
    except (OSError, subprocess.SubprocessError):
        return None
    if out.returncode != 0 or not out.stdout:
        return None

    lines = out.stdout.splitlines()
    if not lines:
        return None
    header = lines[0].split()
    try:
        i_name = header.index("Name")
        i_net = header.index("Network")
        i_in = header.index("Ibytes")
        i_out = header.index("Obytes")
    except ValueError:
        # A Linux netstat has none of these columns; say so rather than guess.
        return None

    counters = {}
    for line in lines[1:]:
        fields = line.split()
        if len(fields) <= max(i_in, i_out, i_net):
            continue
        name = fields[i_name]
        if iface and name != iface:
            continue
        # The link-layer row is the one that counts every packet on the interface;
        # the address rows repeat a subset and would double-count.
        if not fields[i_net].startswith("<Link"):
            continue
        try:
            counters[name] = {
                "bytes_in": int(fields[i_in]),
                "bytes_out": int(fields[i_out]),
            }
        except ValueError:
            continue
    return counters or None


def link_delta(before, after):
    if not before or not after:
        return None
    delta = {}
    for name, start in before.items():
        end = after.get(name)
        if not end:
            continue
        delta[name] = {
            "bytes_in": end["bytes_in"] - start["bytes_in"],
            "bytes_out": end["bytes_out"] - start["bytes_out"],
        }
    return delta or None


# ------------------------------------------------------------------- HTTP helpers

def post_json(url: str, body: dict, headers: dict, timeout: float):
    data = json.dumps(body).encode("utf-8")
    request = urllib.request.Request(url, data=data, headers=headers, method="POST")
    return urllib.request.urlopen(request, timeout=timeout)  # noqa: S310


def discover_model(base_url: str, headers: dict, timeout: float):
    url = base_url.rstrip("/") + "/models"
    try:
        request = urllib.request.Request(url, headers=headers)
        with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310
            payload = json.loads(response.read().decode("utf-8"))
    except (urllib.error.URLError, ValueError, OSError):
        return None
    entries = payload.get("data") or payload.get("models") or []
    if isinstance(entries, list) and entries:
        first = entries[0]
        if isinstance(first, dict):
            return first.get("id") or first.get("name")
        if isinstance(first, str):
            return first
    return None


def build_prompt(target_tokens: int) -> str:
    repeats = max(1, (target_tokens * CHARS_PER_TOKEN) // len(FILLER))
    return (FILLER * repeats) + "\n\n" + QUESTION


# ------------------------------------------------------------------ one streamed run

def stream_once(url, body, headers, timeout):
    """Time to first token, total time, and how many content deltas arrived."""
    started = time.perf_counter()
    first_token_at = None
    tokens = 0
    with post_json(url, body, headers, timeout) as response:
        for raw in response:
            line = raw.decode("utf-8", errors="replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if payload == "[DONE]":
                break
            try:
                chunk = json.loads(payload)
            except ValueError:
                continue
            choices = chunk.get("choices") or []
            if not choices:
                continue
            delta = choices[0].get("delta") or {}
            content = delta.get("content")
            if not content:
                continue
            if first_token_at is None:
                first_token_at = time.perf_counter()
            tokens += 1
    finished = time.perf_counter()
    if first_token_at is None:
        return None
    return {
        "time_to_first_token_s": first_token_at - started,
        "decode_seconds": max(finished - first_token_at, 1e-9),
        "output_tokens": tokens,
        "total_seconds": finished - started,
    }


def bench_once(exo_base, body, headers, timeout):
    """exo's own prefill and decode figures, measured inside the server."""
    url = exo_base.rstrip("/") + "/bench/chat/completions"
    request_body = dict(body)
    request_body["stream"] = False
    with post_json(url, request_body, headers, timeout) as response:
        payload = json.loads(response.read().decode("utf-8"))
    keys = ("prompt_tps", "generation_tps", "prompt_tokens", "generation_tokens",
            "peak_memory_usage")
    return {k: payload[k] for k in keys if k in payload}


def summarise(values):
    if not values:
        return None
    ordered = sorted(values)
    return {
        "median": round(statistics.median(ordered), 4),
        "p90": round(ordered[min(len(ordered) - 1, int(0.9 * len(ordered)))], 4),
        "min": round(ordered[0], 4),
        "max": round(ordered[-1], 4),
        "n": len(ordered),
    }


def main() -> int:
    args = parse_args()

    headers = {"Content-Type": "application/json"}
    key = os.environ.get(args.api_key_env, "")
    if key:
        headers["Authorization"] = f"Bearer {key}"

    model = args.model or discover_model(args.base_url, headers, 30.0)
    if not model:
        print(
            "No model given and the server did not name one at /models. "
            "Pass --model.",
            file=sys.stderr,
        )
        return 1

    prompt = build_prompt(args.prefill_tokens)
    body = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": args.max_tokens,
        "temperature": 0.0,
        "stream": True,
    }
    url = args.base_url.rstrip("/") + "/chat/completions"

    print(f"==> warming up {model}")
    warm = dict(body)
    warm["max_tokens"] = 8
    try:
        stream_once(url, warm, headers, args.timeout)
    except (urllib.error.URLError, OSError) as exc:
        print(f"The server did not answer: {exc}", file=sys.stderr)
        print(f"Checked {url}", file=sys.stderr)
        return 1

    before = read_link_counters(args.iface)
    wall_start = time.perf_counter()

    runs = []
    failures = []
    lock = threading.Lock()
    pending = list(range(args.requests))
    concurrency = max(1, min(args.concurrency, args.requests))

    print(
        f"==> {args.requests} request(s), {args.max_tokens} tokens each, "
        f"{concurrency} in flight"
    )

    def worker():
        while True:
            with lock:
                if not pending:
                    return
                pending.pop()
            try:
                result = stream_once(url, body, headers, args.timeout)
            except (urllib.error.URLError, OSError) as exc:
                with lock:
                    failures.append(str(exc))
                continue
            if result:
                with lock:
                    runs.append(result)
                    print(f"    {len(runs)}/{args.requests}", end="\r", flush=True)

    threads = [threading.Thread(target=worker) for _ in range(concurrency)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    print("")
    for message in failures[:3]:
        print(f"    a request failed: {message}", file=sys.stderr)

    wall_seconds = time.perf_counter() - wall_start
    after = read_link_counters(args.iface)

    if not runs:
        print("No request completed. Nothing to record.", file=sys.stderr)
        return 1

    ttfts = [r["time_to_first_token_s"] for r in runs]
    decode_rates = [
        r["output_tokens"] / r["decode_seconds"] for r in runs if r["output_tokens"]
    ]
    # Prefill rate from time to first token, using the prompt length we asked for.
    prompt_tokens_estimate = len(prompt) // CHARS_PER_TOKEN
    prefill_rates = [prompt_tokens_estimate / t for t in ttfts if t > 0]

    exo_bench = None
    if args.exo_bench_url:
        try:
            exo_bench = bench_once(args.exo_bench_url, body, headers, args.timeout)
        except (urllib.error.URLError, ValueError, OSError) as exc:
            print(f"exo bench endpoint did not answer: {exc}", file=sys.stderr)

    record = {
        "lab": "part-21/two-mac-cluster-over-thunderbolt-5",
        "record": "served-measurement",
        "date": date.today().isoformat(),
        "recorded_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "label": args.label,
        "model": model,
        "base_url": args.base_url,
        "client_machine": platform.node().split(".")[0],
        "client_platform": platform.platform(),
        "requests_completed": len(runs),
        "requests_asked": args.requests,
        "requests_failed": len(failures),
        "concurrency": concurrency,
        "max_tokens": args.max_tokens,
        "prompt_tokens_estimate": prompt_tokens_estimate,
        "wall_seconds": round(wall_seconds, 3),
        "aggregate_output_tokens_per_second": round(
            sum(r["output_tokens"] for r in runs) / wall_seconds, 3
        ),
        "time_to_first_token_s": summarise(ttfts),
        "decode_tokens_per_second": summarise(decode_rates),
        "prefill_tokens_per_second_estimate": summarise(prefill_rates),
        "link_bytes": link_delta(before, after),
        "link_interface": args.iface or "all",
        "exo_bench": exo_bench,
    }

    if record["link_bytes"] is None:
        record["link_note"] = (
            "No BSD netstat interface counters were readable on this machine, so "
            "link use was not recorded. Run this script on one of the Macs."
        )

    text = json.dumps(record, sort_keys=True)
    if args.print_only:
        print(text)
        return 0
    with open(args.labbook, "a", encoding="utf-8") as fh:
        fh.write(text + "\n")
    print(text)
    print(f"Appended to {args.labbook}")
    print("Read the lines back with:")
    print(f"  grep 'served-measurement' {args.labbook}")
    return 0


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