"""Compare a PD (prefill/decode split) deployment against a single-engine server on the same prompts.

    python serve/check_pd.py --ref http://localhost:8888 --pd http://tray10:8888 --turns 3 --prefix-tokens 3000

Greedy, multi-turn with a long shared prefix (so the PE's prefix cache gets exercised); prints per-turn TTFT and
whether the outputs are identical. Exits non-zero on any mismatch.
"""

from __future__ import annotations

import argparse
import json
import time

import requests

p = argparse.ArgumentParser()
p.add_argument("--ref", default="http://localhost:8888")
p.add_argument("--pd", required=True)
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=3000)
p.add_argument("--max-tokens", type=int, default=64)
args = p.parse_args()

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.",
    "Quote the first five words of the document.",
    "What is the last word of the document?",
    "How many distinct word stems are used? Answer briefly.",
]


def ask(url: str, messages: list[dict]) -> tuple[str, float, float, int]:
    body = {
        "model": args.model,
        "messages": messages,
        "stream": True,
        "stream_options": {"include_usage": True},
        "temperature": 0,
        "max_tokens": args.max_tokens,
        "seed": 1,
    }
    t0 = time.perf_counter()
    ttft, text, n = None, [], 0
    with requests.post(f"{url}/v1/chat/completions", 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"):
                n = obj["usage"]["completion_tokens"]
            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, n


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)]})
    ref, ttft_r, tot_r, n_r = ask(args.ref, messages)
    pd, ttft_p, tot_p, n_p = ask(args.pd, messages)
    same = ref == pd
    mismatches += not same
    print(
        f"turn {turn + 1}: ref TTFT {ttft_r * 1e3:6.0f} ms total {tot_r * 1e3:6.0f} ms ({n_r} tok) | PD TTFT {ttft_p * 1e3:6.0f} ms total {tot_p * 1e3:6.0f} ms ({n_p} tok) | "
        f"outputs {'identical' if same else 'DIFFER'}"
    )
    if not same:
        print("   ref:", repr(ref[:160]))
        print("   pd :", repr(pd[:160]))
    messages.append({"role": "assistant", "content": ref})
raise SystemExit(1 if mismatches else 0)