"""
measure_client.py — TTFT / ITL / throughput as the USER sees them.

Works against any OpenAI-compatible streaming server: vLLM, SGLang,
TensorRT-LLM (trtllm-serve), or tiny_server.py from this post.

    vllm serve meta-llama/Llama-3.1-8B-Instruct      # on a GPU box
    python measure_client.py --base-url http://localhost:8000/v1 \
        --model meta-llama/Llama-3.1-8B-Instruct
"""
import argparse
import statistics
import time
from openai import OpenAI

ap = argparse.ArgumentParser()
ap.add_argument("--base-url", default="http://localhost:8000/v1")
ap.add_argument("--model", default="tiny")
ap.add_argument("--max-tokens", type=int, default=64)
ap.add_argument("--prompt", default="Explain how a large language model answers "
                                    "a question, step by step. " * 4)
args = ap.parse_args()

client = OpenAI(base_url=args.base_url, api_key="unused")

t_send = time.perf_counter()
stream = client.completions.create(model=args.model, prompt=args.prompt,
                                   max_tokens=args.max_tokens, temperature=0,
                                   stream=True,
                                   stream_options={"include_usage": True})
arrivals, usage = [], None
for chunk in stream:
    if chunk.choices:                      # one chunk per generated token
        arrivals.append(time.perf_counter())
    if chunk.usage:
        usage = chunk.usage
t_done = time.perf_counter()

ttft = arrivals[0] - t_send
gaps = [b - a for a, b in zip(arrivals, arrivals[1:])]
out_tokens = usage.completion_tokens if usage else len(arrivals)
print(f"prompt tokens      : {usage.prompt_tokens if usage else '?'}")
print(f"output tokens      : {out_tokens}")
print(f"TTFT               : {ttft*1e3:8.1f} ms")
print(f"ITL p50 / p99      : {statistics.median(gaps)*1e3:8.1f} / "
      f"{sorted(gaps)[int(0.99*(len(gaps)-1))]*1e3:.1f} ms")
print(f"end-to-end latency : {t_done - t_send:8.2f} s")
print(f"decode tokens/s    : {(out_tokens - 1) / (arrivals[-1] - arrivals[0]):8.1f}")