#!/usr/bin/env python3
"""Measure whether a server is really reusing the prefix it already computed.

Purpose: send one long shared prefix followed by several different short questions, with
    max_tokens set to 1 so that almost all of the measured time is prefill, and compare
    the first (cold) request against the later (warm) ones. Where the server exposes
    prefix-cache counters on a Prometheus endpoint, the script reports what the run added
    to them, so the timing and the counters can be read against each other. A cache that
    is on but never hit looks exactly like a cache that is off, and only these two
    measurements together tell them apart.
Platform: all (spark, strix, mac, nvidia). Pure Python standard library: no pip install.
Minimum memory: 8 GB on the machine running this script; the floor is set by the server.
Assumes: a server listening and answering POST /v1/chat/completions over plain HTTP (a
    localhost or LAN tool; it does not speak TLS). Run it twice, once with the engine's
    prefix cache enabled and once with it disabled, and compare the two notebook lines.
    Results are appended to the lab notebook as one JSON line per run.

Usage:
    python3 measure-prefix-cache.py --base-url http://127.0.0.1:8000/v1 \\
        --model local-chat --label vllm-apc-on --labbook labbook.md

    python3 measure-prefix-cache.py --base-url http://127.0.0.1:8080/v1 \\
        --model local-chat --prefix-words 1200 --questions 6 \\
        --label llama-server-cache-prompt --labbook labbook.md

    # a prefix of your own: a document, a system prompt, a tool schema
    python3 measure-prefix-cache.py --base-url http://127.0.0.1:8000/v1 \\
        --model local-chat --prefix-file manual.txt --label rag-manual

An API key, if the 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.
"""

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

# The synthetic prefix is deliberately dull and self-similar: what is being measured is
# how long the engine takes to read tokens it has already read, not what it makes of them.
PREFIX_SENTENCES = [
    "The service runs one open-weight language model on a single machine.",
    "Memory is divided between the model weights, the key-value cache and the engine's working buffers.",
    "Prefill reads the prompt and is compute-bound; decode writes the answer and is bandwidth-bound.",
    "The key-value cache is charged per token and comes out of a pool reserved at start-up.",
    "Every measurement is recorded with the hardware, the engine version and the date.",
    "A benchmark without its context is a number somebody will misquote.",
    "Concurrency is the size of the pool divided by the context length each conversation is allowed.",
    "The operator cares about time to first token, time per output token and goodput.",
]

QUESTIONS = [
    "Reply with the single word: one.",
    "Reply with the single word: two.",
    "Reply with the single word: three.",
    "Reply with the single word: four.",
    "Reply with the single word: five.",
    "Reply with the single word: six.",
    "Reply with the single word: seven.",
    "Reply with the single word: eight.",
]


def build_prefix(words: int, path: str | None) -> str:
    """A file's contents, or a synthetic block of roughly `words` words."""
    if path:
        with open(path, encoding="utf-8") as handle:
            text = handle.read().strip()
        if not text:
            raise SystemExit(f"{path} is empty")
        return text
    out: list[str] = []
    count = 0
    i = 0
    while count < words:
        sentence = PREFIX_SENTENCES[i % len(PREFIX_SENTENCES)]
        out.append(f"[{i:04d}] {sentence}")
        count += len(sentence.split()) + 1
        i += 1
    return "\n".join(out)


# ------------------------------------------------------------------------ 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 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 only
            return response.read().decode("utf-8", "replace")
    except (urllib.error.URLError, OSError):
        return None


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

# Generic on purpose: engines rename metrics between releases, so this reports whatever
# the server publishes whose name mentions a cache, rather than asserting that a
# particular counter exists. Nothing matched is itself a result worth recording.
CACHE_METRIC = re.compile(r"(prefix_cache|cache_hit|cache_quer|radix|kv_cache)", 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) -> dict[str, float]:
    """Reads a Prometheus exposition page and keeps the cache-related samples."""
    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 CACHE_METRIC.search(match.group("name")):
            continue
        try:
            value = float(match.group("value"))
        except ValueError:
            continue
        out[match.group("name") + (match.group("labels") or "")] = value
    return out


def metrics_delta(before: dict[str, float], after: dict[str, float]) -> dict[str, float]:
    """What this run 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 one_request(cfg: dict, prefix: str, question: str) -> dict:
    """One prefix-plus-question request with max_tokens 1, timed end to end.

    max_tokens is 1 so that the measurement is dominated by prefill. It does not
    isolate prefill perfectly - the response still travels back over HTTP and one
    decode step still happens - but the constant is the same for every request in the
    run, so the cold-to-warm difference is the cache and not the overhead.
    """
    payload = {
        "model": cfg["model"],
        "messages": [
            {"role": "system", "content": prefix},
            {"role": "user", "content": question},
        ],
        "max_tokens": 1,
        "temperature": 0.0,
        "stream": False,
    }
    started = time.perf_counter()
    body = post_json(cfg["chat_url"], payload, cfg["api_key"], cfg["timeout"])
    elapsed = time.perf_counter() - started
    usage = body.get("usage") or {}
    return {
        "elapsed_s": elapsed,
        "prompt_tokens": usage.get("prompt_tokens"),
        "timings": body.get("timings") if isinstance(body.get("timings"), dict) else None,
    }


def run(cfg: dict, prefix: str, questions: list[str]) -> dict:
    """Cold request, then the warm ones, with the cache counters read either side."""
    before = scrape_metrics(cfg["metrics_url"], cfg["api_key"], cfg["timeout"])

    try:
        cold = one_request(cfg, prefix, questions[0])
    except ServerError as exc:
        return {"error": str(exc)}

    warm: list[dict] = []
    errors: list[str] = []
    for question in questions[1:]:
        try:
            warm.append(one_request(cfg, prefix, question))
        except ServerError as exc:
            errors.append(str(exc))

    after = scrape_metrics(cfg["metrics_url"], cfg["api_key"], cfg["timeout"])
    warm_times = [w["elapsed_s"] for w in warm]

    return {
        "prefix_chars": len(prefix),
        "prompt_tokens": cold["prompt_tokens"],
        "cold_s": round(cold["elapsed_s"], 4),
        "warm_s": {
            "count": len(warm_times),
            "p50": round(statistics.median(warm_times), 4) if warm_times else None,
            "mean": round(statistics.fmean(warm_times), 4) if warm_times else None,
            "min": round(min(warm_times), 4) if warm_times else None,
            "max": round(max(warm_times), 4) if warm_times else None,
        },
        "warm_over_cold": (
            round(statistics.median(warm_times) / cold["elapsed_s"], 4)
            if warm_times and cold["elapsed_s"] > 0 else None
        ),
        "server_timings_sample": cold["timings"],
        "cache_metrics_delta": metrics_delta(before, after),
        "failed": len(errors),
        "first_error": errors[0] if errors else None,
    }


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 parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base-url", default="http://127.0.0.1:8000/v1",
                        help="OpenAI-compatible base URL, ending in /v1")
    parser.add_argument("--model", required=True, help="model name the server reports at /v1/models")
    parser.add_argument("--prefix-words", type=int, default=800,
                        help="approximate length of the synthetic shared prefix, in words")
    parser.add_argument("--prefix-file", default=None,
                        help="use this file as the shared prefix instead of the synthetic one")
    parser.add_argument("--questions", type=int, default=6,
                        help="requests to send: the first is cold, the rest should hit the cache")
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--label", default="prefix-cache", help="tag written into the notebook line")
    parser.add_argument("--engine", default="unknown", help="engine name recorded in the notebook")
    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 server needs one")
    parser.add_argument("--labbook", default="labbook.md")
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)

    parts = urlsplit(args.base_url)
    if parts.scheme != "http":
        print("This tool speaks plain HTTP only; point it at a localhost or LAN endpoint.", file=sys.stderr)
        return 2
    if not parts.hostname:
        print(f"Could not read a host from --base-url {args.base_url!r}.", file=sys.stderr)
        return 2

    root = f"{parts.scheme}://{parts.netloc}"
    path = parts.path.rstrip("/") or "/v1"
    cfg = {
        "chat_url": f"{root}{path}/chat/completions",
        "metrics_url": f"{root}/metrics",
        "model": args.model,
        "timeout": args.timeout,
        "api_key": os.environ.get(args.api_key_env, ""),
    }

    if args.questions < 2:
        print("--questions must be at least 2: one cold request and one warm one.", file=sys.stderr)
        return 2
    questions = [QUESTIONS[i % len(QUESTIONS)] for i in range(args.questions)]
    prefix = build_prefix(args.prefix_words, args.prefix_file)

    print(f"==> {args.label}: {args.model} at {args.base_url}")
    print(f"    shared prefix of {len(prefix)} characters, {args.questions} request(s), max_tokens 1")

    result = run(cfg, prefix, questions)
    if "error" in result:
        print(f"    the first request failed: {result['error']}", file=sys.stderr)
        return 1

    print(f"    prompt      {result['prompt_tokens']} tokens (as counted by the server)")
    print(f"    cold        {result['cold_s']:.3f} s")
    if result["warm_s"]["p50"] is not None:
        print(f"    warm        p50 {result['warm_s']['p50']:.3f} s over {result['warm_s']['count']} request(s)")
        print(f"    ratio       warm/cold {result['warm_over_cold']:.3f}")
    if result["cache_metrics_delta"]:
        print("    cache counters this run added:")
        for name, value in result["cache_metrics_delta"].items():
            print(f"      {name} += {value}")
    else:
        print("    cache       the server exposed no cache counters; the timing is your only evidence")
    if result["first_error"]:
        print(f"    first error {result['first_error']}")

    record = dict(result)
    record.update({
        "lab": "part-17/measure-prefix-cache",
        "label": args.label,
        "engine": args.engine,
        "engine_version": args.engine_version,
        "model": args.model,
        "base_url": args.base_url,
        "prefix_source": args.prefix_file or f"synthetic-{args.prefix_words}-words",
        "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
    })
    append_labbook(args.labbook, record)
    print(f"    appended 1 line to {args.labbook}")
    return 0


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