#!/usr/bin/env python3
"""Measure what speculative decoding does to one server, and prove it changed nothing else.

Purpose: run one fixed prompt set through an OpenAI-compatible endpoint with speculation
    off and again with speculation on, and report the three things that decide whether
    speculation earns its place: output tokens per second, whatever acceptance statistics
    the server exposes on its Prometheus endpoint, and a byte-for-byte comparison of the
    greedy completions. Speculative decoding is a latency technique, not a quality trade,
    so a run whose text changed is a run to investigate rather than to publish.
Platform: all (spark, strix, mac, nvidia). Pure Python standard library: no pip install.
Minimum memory: 8 GB on the machine running this script; the memory floor is set by the
    server, not by the generator, and the generator may run on another machine on the LAN.
Assumes: one or two servers already listening and answering POST /v1/chat/completions over
    plain HTTP (this is a localhost or LAN tool and does not speak TLS). Both endpoints
    must be serving the same weights at the same quantisation, differing only in whether
    speculation is enabled, or the comparison means nothing. Results are appended to the
    lab notebook as one JSON line per endpoint plus one for the comparison.

Usage:
    # measure a single endpoint, e.g. before you have a draft at all
    python3 measure-speculative.py --baseline-url http://127.0.0.1:8080/v1 \\
        --model local-chat --label llama-server-no-draft --labbook labbook.md

    # compare two endpoints: one without a draft, one with
    python3 measure-speculative.py \\
        --baseline-url http://127.0.0.1:8080/v1 \\
        --speculative-url http://127.0.0.1:8081/v1 \\
        --model local-chat --max-tokens 192 --repeats 2 \\
        --label qwen3-8b-draft-0.6b --labbook labbook.md

    # your own prompts, one per line, instead of the built-in mixed set
    python3 measure-speculative.py --baseline-url http://127.0.0.1:8080/v1 \\
        --model local-chat --prompts my-prompts.txt --prompt-set file

An API key, if a server needs one, is read from the environment variable named by
--api-key-env. No key is written to this file, to the output or to the notebook.

Counting note: token counts come from the "usage" object the server returns, so they are
    the server's own count rather than an estimate. A server that omits "usage" is
    reported with a null token count and its tokens-per-second figure is left out rather
    than guessed.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import statistics
import sys
import time
import urllib.error
import urllib.request
from urllib.parse import urlsplit

# --------------------------------------------------------------------------- prompts

# Three shapes of work, because speculation behaves completely differently on each.
# "copy" repeats material that is already in the prompt, which is where lookup-based
# drafting wins; "closed" has one predictable answer; "open" is high-entropy writing,
# where a draft agrees with the target least often.
BUILTIN_PROMPTS = [
    # copy-heavy
    "Here is a function:\n\ndef total(rows):\n    out = 0\n    for r in rows:\n        out += r['amount']\n    return out\n\nRewrite it to skip rows whose 'amount' key is missing, and show the whole function.",
    "Reformat this list as a JSON array of objects with keys name and port: "
    "llama-server 8080, vllm 8000, sglang 30000, litellm 4000.",
    "Copy this sentence exactly, then explain it in one sentence: "
    "decode is bandwidth-bound because every active weight is read once per token.",
    "Take this shell line and add error handling, showing the complete result: "
    "curl -s http://127.0.0.1:8080/v1/models | jq .data",
    # closed
    "In two sentences, what does a KV cache hold and why does it grow with the conversation?",
    "Name the two halves of a generation step and say which one is compute-bound.",
    "What is an acceptance rate in speculative decoding? Answer in one sentence.",
    "List four fields that must accompany a tokens-per-second measurement.",
    # open
    "Write a short paragraph, in your own words, about why a smaller model that fits in memory can beat a larger one that does not.",
    "Describe, as if to a colleague over coffee, what surprised you most about running language models on your own hardware.",
    "Invent a plausible name and one-line description for a tool that records benchmark context automatically.",
    "Write three sentences of encouragement for somebody whose first fine-tune made their model worse.",
]


def load_prompts(prompt_set: str, path: str | None) -> list[str]:
    """The built-in mixed set, or one prompt per non-blank line of a file."""
    if prompt_set == "builtin":
        return list(BUILTIN_PROMPTS)
    if not path:
        raise SystemExit("--prompt-set file needs --prompts pointing at a file")
    with open(path, encoding="utf-8") as handle:
        prompts = [line.strip() for line in handle if line.strip()]
    if not prompts:
        raise SystemExit(f"{path} contained no prompts")
    return prompts


# ------------------------------------------------------------------------ HTTP


class ServerError(Exception):
    """The server did not answer in a way this script can use."""


def post_json(url: str, payload: dict, api_key: str, timeout: float) -> dict:
    """One non-streaming POST. Returns the decoded JSON body."""
    body = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(url, data=body, method="POST")
    request.add_header("Content-Type", "application/json")
    if api_key:
        request.add_header("Authorization", f"Bearer {api_key}")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310 - http(s) only, checked in main
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:200]
        raise ServerError(f"HTTP {exc.code} from {url}: {detail}") from exc
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as exc:
        raise ServerError(f"{type(exc).__name__} talking to {url}: {exc}") from exc


def get_text(url: str, api_key: str, timeout: float) -> str | None:
    """GET a text body, or None if the endpoint is absent or refuses."""
    request = urllib.request.Request(url, method="GET")
    if api_key:
        request.add_header("Authorization", f"Bearer {api_key}")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310 - http(s) only
            return response.read().decode("utf-8", "replace")
    except (urllib.error.URLError, OSError):
        return None


# ------------------------------------------------------------- Prometheus scraping

# Deliberately generic. Engines rename their metrics between releases, so this script
# reports whatever the server publishes whose name mentions speculation or the draft,
# rather than asserting that a particular metric name exists. An empty result means
# "this server exposed nothing under those names", which the page tells you to record.
SPEC_METRIC = re.compile(r"(spec_decode|speculat|draft|accept)", re.IGNORECASE)
SAMPLE = re.compile(r"^(?P<name>[A-Za-z_:][A-Za-z0-9_:]*)(?P<labels>\{[^}]*\})?\s+(?P<value>[-+0-9.eE]+|NaN)$")


def scrape_metrics(metrics_url: str, api_key: str, timeout: float, keep: re.Pattern[str]) -> dict[str, float]:
    """Reads a Prometheus exposition page and keeps the samples whose name matches."""
    text = get_text(metrics_url, api_key, timeout)
    if text is None:
        return {}
    out: dict[str, float] = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        match = SAMPLE.match(line)
        if not match or not keep.search(match.group("name")):
            continue
        try:
            value = float(match.group("value"))
        except ValueError:
            continue
        key = match.group("name") + (match.group("labels") or "")
        out[key] = value
    return out


def metrics_delta(before: dict[str, float], after: dict[str, float]) -> dict[str, float]:
    """What the run itself added, so a warm server's history does not pollute the record."""
    keys = set(before) | set(after)
    return {k: round(after.get(k, 0.0) - before.get(k, 0.0), 4) for k in sorted(keys)}


# ------------------------------------------------------------------------ measurement


def run_endpoint(cfg: dict, prompts: list[str], repeats: int) -> dict:
    """Sends every prompt `repeats` times, sequentially, and times each request.

    Sequential and single-stream on purpose: speculation spends idle arithmetic
    capacity, so concurrency one is where it has the most to gain, and it is the
    condition under which the acceptance arithmetic on the lesson page applies.
    """
    completions: list[str] = []
    latencies: list[float] = []
    token_counts: list[int | None] = []
    timings: list[dict] = []
    errors: list[str] = []

    metrics_before = scrape_metrics(cfg["metrics_url"], cfg["api_key"], cfg["timeout"], SPEC_METRIC)
    started = time.perf_counter()

    for _ in range(repeats):
        for prompt in prompts:
            payload = {
                "model": cfg["model"],
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": cfg["max_tokens"],
                "temperature": 0.0,
                "top_p": 1.0,
                "seed": cfg["seed"],
                "stream": False,
            }
            t0 = time.perf_counter()
            try:
                body = post_json(cfg["chat_url"], payload, cfg["api_key"], cfg["timeout"])
            except ServerError as exc:
                errors.append(str(exc))
                continue
            latencies.append(time.perf_counter() - t0)
            choices = body.get("choices") or [{}]
            message = choices[0].get("message") or {}
            completions.append(message.get("content") or "")
            usage = body.get("usage") or {}
            completion_tokens = usage.get("completion_tokens")
            token_counts.append(int(completion_tokens) if isinstance(completion_tokens, int) else None)
            if isinstance(body.get("timings"), dict):
                timings.append(body["timings"])

    wall = time.perf_counter() - started
    metrics_after = scrape_metrics(cfg["metrics_url"], cfg["api_key"], cfg["timeout"], SPEC_METRIC)

    counted = [t for t in token_counts if t is not None]
    total_tokens = sum(counted) if counted else None
    complete = len(counted) == len(latencies) and bool(counted)

    return {
        "label": cfg["label"],
        "base_url": cfg["base_url"],
        "model": cfg["model"],
        "requests": len(prompts) * repeats,
        "completed": len(latencies),
        "failed": len(errors),
        "first_error": errors[0] if errors else None,
        "wall_s": round(wall, 3),
        "latency_s": {
            "p50": round(statistics.median(latencies), 4) if latencies else None,
            "mean": round(statistics.fmean(latencies), 4) if latencies else None,
            "max": round(max(latencies), 4) if latencies else None,
        },
        "output_tokens": total_tokens,
        "output_tokens_per_s": (
            round(total_tokens / wall, 2) if complete and total_tokens and wall > 0 else None
        ),
        "tokens_counted_by_server": complete,
        "server_timings_sample": timings[0] if timings else None,
        "spec_metrics_delta": metrics_delta(metrics_before, metrics_after),
        "completions": completions,
    }


def compare_text(baseline: list[str], speculative: list[str]) -> dict:
    """Byte-for-byte comparison of the two sets of greedy completions.

    Verification means the accepted tokens are the ones the target model would have
    produced, so at temperature zero the two runs should agree exactly. They can still
    differ for reasons that have nothing to do with the draft - a different batch
    composition changes the order of floating-point reductions, and the engines say so -
    which is why this reports where the first divergence is rather than only that there
    was one.
    """
    pairs = list(zip(baseline, speculative))
    identical = [i for i, (a, b) in enumerate(pairs) if a == b]
    divergent = []
    for i, (a, b) in enumerate(pairs):
        if a == b:
            continue
        cut = 0
        for cut, (ca, cb) in enumerate(zip(a, b)):
            if ca != cb:
                break
        else:
            cut = min(len(a), len(b))
        divergent.append(
            {
                "prompt_index": i,
                "first_difference_at_char": cut,
                "baseline_tail": a[cut : cut + 60],
                "speculative_tail": b[cut : cut + 60],
            }
        )
    return {
        "compared": len(pairs),
        "identical": len(identical),
        "divergent": len(divergent),
        "divergences": divergent[:5],
    }


# ---------------------------------------------------------------------------- output


def print_endpoint(row: dict) -> None:
    """One endpoint's result, as an aligned block."""
    print(f"  {row['label']}  ({row['base_url']})")
    print(f"    requests    {row['completed']}/{row['requests']} completed, {row['failed']} failed")
    print(f"    wall clock  {row['wall_s']:.2f} s")
    if row["output_tokens_per_s"] is not None:
        print(f"    output      {row['output_tokens']} tokens, {row['output_tokens_per_s']:.2f} tokens/s")
    else:
        print("    output      the server did not return a usage object; tokens/s not computed")
    if row["latency_s"]["p50"] is not None:
        print(f"    latency     p50 {row['latency_s']['p50']:.3f} s, max {row['latency_s']['max']:.3f} s")
    if row["spec_metrics_delta"]:
        print("    acceptance statistics from the server's /metrics during this run:")
        for name, value in row["spec_metrics_delta"].items():
            print(f"      {name} += {value}")
    else:
        print("    acceptance  the server exposed no speculation metrics; record that as the finding")
    if row["first_error"]:
        print(f"    first error {row['first_error']}")


def append_labbook(path: str, record: dict) -> None:
    """Appends one JSON line, in the course notebook format."""
    with open(path, "a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")


def endpoint_config(base_url: str, args: argparse.Namespace, label: str) -> dict:
    """Derives the chat and metrics URLs from an OpenAI-compatible base URL."""
    parts = urlsplit(base_url)
    if parts.scheme != "http":
        raise SystemExit("This tool speaks plain HTTP only; point it at a localhost or LAN endpoint.")
    if not parts.hostname:
        raise SystemExit(f"Could not read a host from {base_url!r}.")
    root = f"{parts.scheme}://{parts.netloc}"
    path = parts.path.rstrip("/") or "/v1"
    return {
        "base_url": base_url,
        "chat_url": f"{root}{path}/chat/completions",
        "metrics_url": f"{root}/metrics",
        "model": args.model,
        "max_tokens": args.max_tokens,
        "seed": args.seed,
        "timeout": args.timeout,
        "api_key": os.environ.get(args.api_key_env, ""),
        "label": label,
    }


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--baseline-url", required=True,
                        help="OpenAI-compatible base URL of the server WITHOUT speculation, ending in /v1")
    parser.add_argument("--speculative-url", default=None,
                        help="base URL of the server WITH speculation; omit to measure one endpoint only")
    parser.add_argument("--model", required=True, help="model name the server reports at /v1/models")
    parser.add_argument("--max-tokens", type=int, default=192)
    parser.add_argument("--seed", type=int, default=0,
                        help="sent as the request seed; ignored by servers that do not accept one")
    parser.add_argument("--repeats", type=int, default=1, help="times to send the whole prompt set")
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--prompt-set", default="builtin", choices=["builtin", "file"])
    parser.add_argument("--prompts", default=None, help="file of prompts, one per line, for --prompt-set file")
    parser.add_argument("--label", default="speculative-run", help="tag written into the notebook lines")
    parser.add_argument("--engine", default="unknown", help="engine name recorded in the notebook, e.g. llama.cpp")
    parser.add_argument("--engine-version", default="unknown", help="engine version recorded in the notebook")
    parser.add_argument("--api-key-env", default="SPEC_API_KEY",
                        help="environment variable holding the API key, if the servers need one")
    parser.add_argument("--labbook", default="labbook.md")
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    prompts = load_prompts(args.prompt_set, args.prompts)

    print(f"==> {args.label}: {args.model}, {len(prompts)} prompt(s) x {args.repeats} repeat(s), "
          f"max_tokens {args.max_tokens}, temperature 0")

    baseline = run_endpoint(endpoint_config(args.baseline_url, args, f"{args.label}/baseline"), prompts, args.repeats)
    print_endpoint(baseline)

    speculative = None
    comparison = None
    if args.speculative_url:
        speculative = run_endpoint(
            endpoint_config(args.speculative_url, args, f"{args.label}/speculative"), prompts, args.repeats
        )
        print_endpoint(speculative)
        comparison = compare_text(baseline["completions"], speculative["completions"])
        print("  comparison")
        print(f"    text        {comparison['identical']}/{comparison['compared']} completions identical, "
              f"{comparison['divergent']} divergent")
        for d in comparison["divergences"]:
            print(f"      prompt {d['prompt_index']}: first difference at character {d['first_difference_at_char']}")
        if baseline["output_tokens_per_s"] and speculative["output_tokens_per_s"]:
            ratio = speculative["output_tokens_per_s"] / baseline["output_tokens_per_s"]
            print(f"    speed       {ratio:.2f}x the baseline token rate on this prompt set")

    stamp = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    for row in (baseline, speculative):
        if row is None:
            continue
        record = dict(row)
        # The completions are the evidence for the comparison, not part of the record:
        # a notebook line has to stay readable.
        record.pop("completions", None)
        record.update({
            "lab": "part-17/measure-speculative",
            "engine": args.engine,
            "engine_version": args.engine_version,
            "max_tokens": args.max_tokens,
            "repeats": args.repeats,
            "prompt_count": len(prompts),
            "recorded_at": stamp,
        })
        append_labbook(args.labbook, record)

    if comparison is not None:
        append_labbook(args.labbook, {
            "lab": "part-17/measure-speculative",
            "label": f"{args.label}/comparison",
            "engine": args.engine,
            "engine_version": args.engine_version,
            "model": args.model,
            "baseline_tokens_per_s": baseline["output_tokens_per_s"],
            "speculative_tokens_per_s": speculative["output_tokens_per_s"] if speculative else None,
            "text_comparison": comparison,
            "recorded_at": stamp,
        })

    print(f"    appended {2 if comparison else 1} line(s) to {args.labbook}")

    if baseline["completed"] == 0 or (speculative is not None and speculative["completed"] == 0):
        print("    at least one endpoint completed no requests; see the error above.", file=sys.stderr)
        return 1
    if comparison is not None and comparison["divergent"]:
        print("    completions differed between the two servers; the page explains what to check.",
              file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))