"""A/B check of the host prefix cache on a running serve/openai_server.py.

Builds a multi-turn conversation with a long shared prefix, sends every turn twice (cache bypassed via
"prefix_cache": false, then with the cache) with greedy sampling, and compares outputs and TTFT.

    python serve/check_prefix_cache.py --port 8888 --turns 3 --prefix-tokens 6000
"""

from __future__ import annotations

import argparse
import json
import time

import requests

p = argparse.ArgumentParser()
p.add_argument("--port", type=int, default=8888)
p.add_argument("--model", default="deepseek-ai/DeepSeek-V4.1-Flash-1017")
p.add_argument("--turns", type=int, default=3)
p.add_argument("--prefix-tokens", type=int, default=6000, help="approximate size of the shared document")
p.add_argument("--max-tokens", type=int, default=64)
args = p.parse_args()
URL = f"http://localhost:{args.port}/v1/chat/completions"

# A deterministic pseudo-document; ~1 token per word for this tokenizer is close enough.
words = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa", "lambda", "mu"]
doc = " ".join(f"{words[i % len(words)]}{i % 97}" for i in range(args.prefix_tokens))
questions = [
    "Which word appears right after 'gamma2'? Answer briefly.",
    "How many distinct word stems are used in the document? Answer briefly.",
    "Quote the first five words of the document.",
    "What is the last word of the document?",
]


def ask(messages: list[dict], use_cache: bool) -> tuple[str, float, float, dict]:
    body = {
        "model": args.model,
        "messages": messages,
        "stream": True,
        "stream_options": {"include_usage": True},
        "temperature": 0,
        "max_tokens": args.max_tokens,
        "seed": 1,
        "prefix_cache": use_cache,
    }
    t0 = time.perf_counter()
    ttft = None
    text, usage = [], {}
    with requests.post(URL, json=body, stream=True, timeout=1800) as r:
        r.raise_for_status()
        for line in r.iter_lines():
            if not line or not line.startswith(b"data: ") or line == b"data: [DONE]":
                continue
            obj = json.loads(line[6:])
            if obj.get("usage"):
                usage = obj["usage"]
            for ch in obj.get("choices", []):
                delta = ch["delta"].get("content")
                if delta:
                    if ttft is None:
                        ttft = time.perf_counter() - t0
                    text.append(delta)
    return "".join(text), ttft or 0.0, time.perf_counter() - t0, usage


messages = [{"role": "system", "content": "You answer questions about the document below.\n\n" + doc}]
mismatches = 0
for turn in range(args.turns):
    messages.append({"role": "user", "content": questions[turn % len(questions)]})
    base, ttft_b, tot_b, usage_b = ask(messages, use_cache=False)
    cached, ttft_c, tot_c, usage_c = ask(messages, use_cache=True)
    same = base == cached
    mismatches += not same
    print(
        f"turn {turn + 1}: prompt {usage_c.get('prompt_tokens')} tok | TTFT no-cache {ttft_b * 1e3:7.0f} ms, cache {ttft_c * 1e3:7.0f} ms "
        f"({ttft_b / max(ttft_c, 1e-6):.1f}x) | outputs {'identical' if same else 'DIFFER'}"
    )
    if not same:
        print("   no-cache:", repr(base[:200]))
        print("   cache   :", repr(cached[:200]))
    messages.append({"role": "assistant", "content": base})

stats = requests.get(f"http://localhost:{args.port}/deepinfer/metrics", timeout=30).text
print("\n".join(line for line in stats.splitlines() if "hicache:hits" in line or "hicache:lookups" in line or "hit_rate" in line))
raise SystemExit(1 if mismatches else 0)