#!/usr/bin/env python3
"""Measure key-value reuse across the turns of a conversation, and across two engines.

Purpose: send a growing multi-turn conversation to one or two OpenAI-compatible endpoints,
    time to the first token of each turn, and read each engine's prefix-cache counters
    either side of the run. With one endpoint this measures reuse across turns, which is
    what an agent loop or a chat client does. With two endpoints given, turns alternate
    between them, which measures whether a shared cache tier is letting one engine reuse
    what the other computed. Appends one notebook line per turn plus one summary line.
Platform: all (spark, strix, mac, nvidia). Pure Python standard library: no pip install.
Minimum memory: 1 GB for the tool itself; 16 GB on the machine serving the model.
Assumes: Python 3.9 or later; a server answering POST /v1/chat/completions with streaming
    over plain HTTP; the model name the server reports at /v1/models. Counters are read
    from /metrics where the engine exposes it, which vLLM does by default and llama-server
    does when started with --metrics. An API key, if the server needs one, is read from
    the environment variable named by --api-key-env and is never written to the notebook.

Usage:
    python3 measure-reuse.py --base-url http://127.0.0.1:8000/v1 --model local-chat \\
        --turns 6 --preamble-words 1500 --label offload-on --labbook labbook.md

    python3 measure-reuse.py \\
        --base-url http://127.0.0.1:8100/v1 \\
        --base-url http://127.0.0.1:8200/v1 \\
        --model local-chat --turns 6 --label shared-store --labbook labbook.md

Counting note: time to first token is measured from just before the request is written to
    the moment the first streamed chunk carrying content arrives. It therefore includes
    connection setup, which is small on a LAN and is the same for every turn, so it does
    not distort the comparison between turns.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from http.client import HTTPResponse
from urllib.parse import urlsplit

PREAMBLE_SENTENCES = [
    "The operator runs one open-weight language model on machines they own.",
    "Prefill reads the prompt and is limited by arithmetic throughput.",
    "Decode writes the answer one token at a time and is limited by memory bandwidth.",
    "A key-value cache holds the key and value vectors for every token already read.",
    "Its size is layers times key-value heads times head dimension times two tensors.",
    "That figure is multiplied by the bytes per element and by the number of tokens.",
    "A prefix cache reuses those vectors when two prompts begin with the same tokens.",
    "The match must start at position zero and must be exact on token identifiers.",
    "Reuse of that kind is bit-identical, so it cannot change what the model produces.",
    "Moving cold blocks to host memory extends the cache without touching the network.",
    "Writing blocks to disk lets a cache survive a restart of the serving process.",
    "A shared store lets a second engine reuse blocks the first engine computed.",
    "The operator measures time to first token, time per output token and goodput.",
    "Throughput on its own answers a different question from goodput under a deadline.",
]

QUESTIONS = [
    "In one sentence, what limits decode speed?",
    "In one sentence, what does a prefix cache reuse?",
    "In one sentence, why must a prefix match start at position zero?",
    "In one sentence, what does moving blocks to host memory cost?",
    "In one sentence, when is a disk tier worth having?",
    "In one sentence, what does a shared store let a second engine do?",
    "In one sentence, how is a key-value cache size computed?",
    "In one sentence, what is the difference between throughput and goodput?",
    "In one sentence, why does an agent loop benefit from cache reuse?",
    "In one sentence, what does a cache miss look like from the outside?",
]


def build_preamble(words: int) -> str:
    """A stable block of text of roughly the requested length, identical on every run."""
    out: list[str] = []
    count = 0
    i = 0
    while count < words:
        sentence = PREAMBLE_SENTENCES[i % len(PREAMBLE_SENTENCES)]
        out.append(sentence)
        count += len(sentence.split())
        i += 1
    return " ".join(out)


def post_stream(base_url: str, body: dict, api_key: str, timeout: float):
    """Sends one streaming request. Returns (ttft_seconds, total_seconds, chunks, text)."""
    parts = urlsplit(base_url)
    url = f"{base_url.rstrip('/')}/chat/completions"
    if parts.scheme != "http":
        raise ValueError("This tool speaks plain HTTP only; use a localhost or LAN endpoint.")
    data = json.dumps(body).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=data, headers=headers, method="POST")

    started = time.perf_counter()
    ttft = None
    chunks = 0
    pieces: list[str] = []
    response: HTTPResponse
    with urllib.request.urlopen(request, timeout=timeout) as response:
        for raw in response:
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if payload == "[DONE]":
                break
            try:
                obj = json.loads(payload)
            except json.JSONDecodeError:
                continue
            for choice in obj.get("choices", []):
                piece = (choice.get("delta") or {}).get("content")
                if not piece:
                    continue
                if ttft is None:
                    ttft = time.perf_counter() - started
                chunks += 1
                pieces.append(piece)
    total = time.perf_counter() - started
    return ttft, total, chunks, "".join(pieces)


def metrics(base_url: str, timeout: float = 10.0) -> dict:
    """Prefix-cache counters from the engine's /metrics endpoint, where it has one."""
    parts = urlsplit(base_url)
    root = f"{parts.scheme}://{parts.netloc}/metrics"
    wanted = {
        "vllm:prefix_cache_hits": "prefix_cache_hits",
        "vllm:prefix_cache_queries": "prefix_cache_queries",
        "vllm:kv_cache_usage_perc": "kv_cache_usage_perc",
        "llamacpp:prompt_tokens_total": "prompt_tokens_total",
        "llamacpp:tokens_predicted_total": "tokens_predicted_total",
    }
    found: dict[str, float] = {}
    try:
        with urllib.request.urlopen(root, timeout=timeout) as response:
            text = response.read().decode("utf-8", "replace")
    except (urllib.error.URLError, OSError, ValueError):
        return found
    for line in text.splitlines():
        if line.startswith("#") or not line.strip():
            continue
        name, _, rest = line.partition("{")
        if not rest:
            name, _, value = line.partition(" ")
        else:
            _, _, value = rest.partition("} ")
        key = wanted.get(name.strip())
        if not key:
            continue
        try:
            found[key] = found.get(key, 0.0) + float(value.strip())
        except ValueError:
            continue
    return found


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base-url", action="append", default=[],
                        help="OpenAI-compatible base URL ending in /v1; give twice to alternate")
    parser.add_argument("--model", required=True, help="model name the server reports")
    parser.add_argument("--turns", type=int, default=6, help="conversation turns to send")
    parser.add_argument("--preamble-words", type=int, default=1500,
                        help="approximate length of the shared preamble in words")
    parser.add_argument("--max-tokens", type=int, default=48)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--label", default="reuse", help="tag written into the notebook line")
    parser.add_argument("--api-key-env", default="LOADTEST_API_KEY",
                        help="environment variable holding the API key, if one is needed")
    parser.add_argument("--labbook", default="labbook.md", help="notebook to append to")
    parser.add_argument("--lab", default="part-22/lab-kv-cache-offload-and-sharing")
    parser.add_argument("--note", default="", help="free text: tier, connector, anything")
    parser.add_argument("--print-only", action="store_true", help="print it, record nothing")
    args = parser.parse_args()

    urls = args.base_url or ["http://127.0.0.1:8000/v1"]
    if args.turns < 2:
        print("--turns must be at least 2; the point is the comparison between them.",
              file=sys.stderr)
        return 2

    api_key = os.environ.get(args.api_key_env, "")
    preamble = build_preamble(args.preamble_words)

    print(f"==> {args.label}: {args.model}")
    for url in urls:
        print(f"    endpoint {url}")
    print(f"    {args.turns} turns, preamble about {args.preamble_words} words, "
          f"max_tokens {args.max_tokens}")
    if len(urls) > 1:
        print("    turns alternate between the endpoints: a hit on turn two means the")
        print("    second engine reused what the first one computed.")
    print()

    before = {url: metrics(url) for url in urls}

    messages = [{"role": "system", "content": preamble}]
    rows = []
    for turn in range(args.turns):
        url = urls[turn % len(urls)]
        question = QUESTIONS[turn % len(QUESTIONS)]
        messages.append({"role": "user", "content": question})
        body = {
            "model": args.model,
            "messages": messages,
            "max_tokens": args.max_tokens,
            "temperature": args.temperature,
            "stream": True,
        }
        try:
            ttft, total, chunks, text = post_stream(url, body, api_key, args.timeout)
        except (urllib.error.URLError, OSError, ValueError) as exc:
            print(f"    turn {turn + 1}: failed against {url}: {exc}", file=sys.stderr)
            return 1
        messages.append({"role": "assistant", "content": text})
        rows.append({
            "turn": turn + 1,
            "endpoint": url,
            "ttft_s": round(ttft, 4) if ttft is not None else None,
            "total_s": round(total, 4),
            "chunks": chunks,
        })
        shown = f"{ttft:.3f}" if ttft is not None else "  none"
        print(f"    turn {turn + 1:>2}  {url:<34}  TTFT {shown} s  total {total:6.2f} s")

    after = {url: metrics(url) for url in urls}

    print()
    first = rows[0]["ttft_s"]
    later = [r["ttft_s"] for r in rows[1:] if r["ttft_s"] is not None]
    if first and later:
        best = min(later)
        print(f"    turn 1 time to first token   {first:.3f} s")
        print(f"    fastest later turn           {best:.3f} s  ({best / first:.2f}x of turn 1)")
        print("    A later turn much faster than the first is reuse working. A later turn")
        print("    as slow as the first means every turn is being prefilled from scratch,")
        print("    which is what an unhit cache looks like from outside.")

    deltas = {}
    for url in urls:
        b, a = before.get(url, {}), after.get(url, {})
        delta = {k: round(a[k] - b.get(k, 0.0), 3) for k in a if k in ("prefix_cache_hits",
                                                                      "prefix_cache_queries",
                                                                      "prompt_tokens_total",
                                                                      "tokens_predicted_total")}
        if "kv_cache_usage_perc" in a:
            delta["kv_cache_usage_perc_end"] = round(a["kv_cache_usage_perc"], 4)
        if delta:
            deltas[url] = delta
            hits = delta.get("prefix_cache_hits")
            queries = delta.get("prefix_cache_queries")
            if hits is not None and queries:
                print(f"    {url}: prefix cache {hits:.0f} hit tokens of "
                      f"{queries:.0f} queried, {hits / queries * 100:.1f}%")

    if not deltas:
        print("    No /metrics counters were readable. vLLM serves them by default;")
        print("    llama-server needs --metrics. The turn timings above still stand.")

    record = {
        "lab": args.lab,
        "record": "reuse",
        "label": args.label,
        "model": args.model,
        "endpoints": urls,
        "turns": args.turns,
        "preamble_words": args.preamble_words,
        "max_tokens": args.max_tokens,
        "rows": rows,
        "metric_deltas": deltas,
        "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 open(args.labbook, "a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")
    print(f"\n    Appended one reuse line to {args.labbook}")
    return 0


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