#!/usr/bin/env python3
"""Direct OpenAI-compatible streaming benchmark for Qwen3.8-Flash-Next on vLLM.

Measures, per request:
  - exact prompt/completion tokens from server usage
  - client TTFT (monotonic clock)
  - decode tok/s over the post-first-token window
  - overall tok/s
  - server-side prefill tok/s from vLLM metrics counters
  - prefix-cache hits (must stay 0 for cold-prefix validity)

Usage:
  BASE_URL=http://127.0.0.1:18000/v1 MODEL=qwen3.8-flash-next python3 bench_direct.py
"""
import json
import os
import random
import statistics
import time
from concurrent.futures import ThreadPoolExecutor

import requests

BASE = os.environ.get("BASE_URL", "http://127.0.0.1:18000/v1").rstrip("/")
ROOT = BASE[:-3] if BASE.endswith("/v1") else BASE
MODEL = os.environ.get("MODEL", "qwen3.8-flash-next")
OUT_TOKENS = int(os.environ.get("OUT_TOKENS", "256"))
REPEATS = int(os.environ.get("REPEATS", "3"))
INPUTS = [int(x) for x in os.environ.get("INPUTS", "128,512,2048,8192,12000").split(",")]
RAW_PATH = os.environ.get("RAW_PATH", "bench_direct_raw.jsonl")
SUMMARY_PATH = os.environ.get("SUMMARY_PATH", "bench_direct_summary.json")

WORDS = (
    "amber birch cobalt delta ember fjord garnet harbor indigo juniper kelp lumen "
    "meadow nimbus onyx prairie quartz raven saffron tundra umber violet willow xenon "
    "yarrow zephyr alder basalt cedar dune elm flint grove heath isle jet kestrel "
    "larch moss nectar oak pine quill reed slate thistle upland vale wren yew zinnia "
    "canyon brook ridge summit glacier lagoon mesa oasis reef savanna steppe taiga "
    "vista crest delta marsh plain shoal spring trace hollow knoll glen moor"
).split()


def tokenize_count(text: str) -> int:
    r = requests.post(
        ROOT + "/tokenize",
        json={"model": MODEL, "prompt": text, "add_special_tokens": False},
        timeout=120,
    )
    r.raise_for_status()
    return int(r.json()["count"])


def make_prompt(target: int, seed: int):
    """Unique natural-language prompt, tuned to target token count via /tokenize."""
    rng = random.Random(seed)
    nonce = f"request-{rng.getrandbits(64):016x}-{seed}"
    n = max(4, int(target / 1.3))
    text, count = "", 0
    for _ in range(14):
        words = [rng.choice(WORDS) for _ in range(n)]
        text = nonce + "\n" + " ".join(words) + "\n"
        count = tokenize_count(text)
        if abs(count - target) <= max(2, int(target * 0.01)):
            break
        n = max(2, int(n * target / max(count, 1)))
    return text, count


def metrics() -> dict:
    txt = requests.get(ROOT + "/metrics", timeout=60).text
    out = {}
    for line in txt.splitlines():
        if not line or line.startswith("#"):
            continue
        try:
            name = line.split("{", 1)[0].split(" ")[0]
            val = float(line.rsplit(" ", 1)[1])
        except (ValueError, IndexError):
            continue
        out[name] = out.get(name, 0.0) + val
    return out


def delta(after: dict, before: dict, key: str) -> float:
    return after.get(key, 0.0) - before.get(key, 0.0)


def one_request(prompt: str, out_tokens: int, tag: str, seed: int) -> dict:
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": out_tokens,
        "min_tokens": out_tokens,
        "ignore_eos": True,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
        "chat_template_kwargs": {"enable_thinking": False},
    }
    m0 = metrics()
    t0 = time.monotonic()
    resp = requests.post(BASE + "/chat/completions", json=payload, stream=True, timeout=1200)
    resp.raise_for_status()
    t_first = None
    usage = None
    chunks = 0
    for raw in resp.iter_lines(decode_unicode=True):
        if not raw or not raw.startswith("data: "):
            continue
        data = raw[6:]
        if data == "[DONE]":
            break
        obj = json.loads(data)
        if obj.get("usage"):
            usage = obj["usage"]
        for choice in obj.get("choices", []):
            d = choice.get("delta", {})
            piece = d.get("content") or d.get("reasoning_content") or ""
            if piece:
                chunks += 1
                if t_first is None:
                    t_first = time.monotonic()
    t_end = time.monotonic()

    # request metrics are updated at completion; allow a short settle window
    m1 = metrics()
    for _ in range(40):
        if delta(m1, m0, "vllm:request_success_total") >= 1:
            break
        time.sleep(0.05)
        m1 = metrics()

    if t_first is None:
        t_first = t_end
    usage = usage or {}
    pt = int(usage.get("prompt_tokens", 0))
    ct = int(usage.get("completion_tokens", 0))
    ttft = t_first - t0
    decode_window = max(t_end - t_first, 1e-9)
    d_prompt = delta(m1, m0, "vllm:prompt_tokens_total")
    d_prefill_s = delta(m1, m0, "vllm:request_prefill_time_seconds_sum")
    d_prefill_n = delta(m1, m0, "vllm:request_prefill_time_seconds_count")
    d_cache_hits = delta(m1, m0, "vllm:prefix_cache_hits_total")
    d_cache_q = delta(m1, m0, "vllm:prefix_cache_queries_total")
    rec = {
        "tag": tag,
        "seed": seed,
        "input_target": None,
        "prompt_tokens": pt,
        "server_prompt_tokens": d_prompt,
        "completion_tokens": ct,
        "content_chunks": chunks,
        "ttft_s": round(ttft, 4),
        "total_s": round(t_end - t0, 4),
        "decode_s": round(decode_window, 4),
        "decode_tps": round((ct - 1) / decode_window, 2) if ct > 1 else None,
        "overall_tps": round(ct / (t_end - t0), 2) if t_end > t0 else None,
        "server_prefill_s": round(d_prefill_s, 6),
        "server_prefill_tps": round(d_prompt / d_prefill_s, 2) if d_prefill_s > 0 else None,
        "client_prefill_est_tps": round(pt / ttft, 2) if ttft > 0 else None,
        "prefix_cache_hits": d_cache_hits,
        "prefix_cache_queries": d_cache_q,
        "requests_counted": d_prefill_n,
    }
    return rec


def med(xs):
    xs = [x for x in xs if x is not None]
    return round(statistics.median(xs), 2) if xs else None


def rng(xs):
    xs = [x for x in xs if x is not None]
    return [round(min(xs), 2), round(max(xs), 2)] if xs else None


def main():
    raw = []
    print(f"endpoint={BASE} model={MODEL} out={OUT_TOKENS} repeats={REPEATS} inputs={INPUTS}")
    print("cold-start request (first inference after server startup):")
    cold_prompt, cold_count = make_prompt(512, 999999)
    cold = one_request(cold_prompt, OUT_TOKENS, "cold", 999999)
    cold["input_target"] = cold_count
    raw.append(cold)
    print(json.dumps(cold, indent=2))

    for target in INPUTS:
        # warmup must not share a prefix with any measured request
        warm_prompt, warm_count = make_prompt(target, 900000 + target)
        one_request(warm_prompt, OUT_TOKENS, "warmup", 900000 + target)
        for rep in range(REPEATS):
            seed = 1000 + target + rep * 7919
            p, c = make_prompt(target, seed)
            rec = one_request(p, OUT_TOKENS, "measured", seed)
            # prefix caching is enabled for comparability; a non-zero hit means
            # the prompt collided, so retry with a fresh unique prompt
            tries = 0
            while rec["prefix_cache_hits"] > 0 and tries < 3:
                tries += 1
                seed += 104729 * tries
                p, c = make_prompt(target, seed)
                rec = one_request(p, OUT_TOKENS, "measured", seed)
            rec["input_target"] = c
            raw.append(rec)
            print(
                f"in~{target} actual={rec['prompt_tokens']} out={rec['completion_tokens']} "
                f"ttft={rec['ttft_s']:.3f}s decode={rec['decode_tps']} tok/s "
                f"prefill(server)={rec['server_prefill_tps']} tok/s cache_hits={rec['prefix_cache_hits']}"
            )

    # optional small-concurrency aggregate, clearly labelled
    conc_rows = []
    for conc in (2, 4):
        prompts = [make_prompt(2048, 50000 + conc * 100 + i)[0] for i in range(conc)]
        t0 = time.monotonic()
        with ThreadPoolExecutor(max_workers=conc) as ex:
            futs = [ex.submit(one_request, p, OUT_TOKENS, f"conc{conc}", i) for i, p in enumerate(prompts)]
            recs = [f.result() for f in futs]
        wall = time.monotonic() - t0
        total_out = sum(r["completion_tokens"] for r in recs)
        total_in = sum(r["prompt_tokens"] for r in recs)
        row = {
            "concurrency": conc,
            "wall_s": round(wall, 3),
            "requests": conc,
            "total_completion_tokens": total_out,
            "total_prompt_tokens": total_in,
            "aggregate_output_tps": round(total_out / wall, 2),
            "aggregate_prompt_tps": round(total_in / wall, 2),
            "per_stream_decode_tps_median": med([r["decode_tps"] for r in recs]),
            "per_stream_ttft_s_median": med([r["ttft_s"] for r in recs]),
        }
        conc_rows.append(row)
        raw.extend(recs)
        print("concurrency:", json.dumps(row))

    with open(RAW_PATH, "w") as f:
        for rec in raw:
            f.write(json.dumps(rec) + "\n")

    summary = {"endpoint": BASE, "model": MODEL, "out_tokens": OUT_TOKENS,
               "repeats": REPEATS, "inputs": INPUTS, "cold": cold,
               "by_input": [], "concurrency": conc_rows}
    for target in INPUTS:
        rows = [r for r in raw if r["tag"] == "measured" and r["input_target"] is not None
                and abs(r["input_target"] - target) <= max(4, target * 0.02)]
        if not rows:
            continue
        summary["by_input"].append({
            "target": target,
            "actual_prompt_tokens_median": med([r["prompt_tokens"] for r in rows]),
            "n": len(rows),
            "ttft_s_median": med([r["ttft_s"] for r in rows]),
            "ttft_s_range": rng([r["ttft_s"] for r in rows]),
            "decode_tps_median": med([r["decode_tps"] for r in rows]),
            "decode_tps_range": rng([r["decode_tps"] for r in rows]),
            "server_prefill_tps_median": med([r["server_prefill_tps"] for r in rows]),
            "server_prefill_tps_range": rng([r["server_prefill_tps"] for r in rows]),
            "client_prefill_est_tps_median": med([r["client_prefill_est_tps"] for r in rows]),
            "overall_tps_median": med([r["overall_tps"] for r in rows]),
            "prefix_cache_hits_total": sum(r["prefix_cache_hits"] for r in rows),
        })
    with open(SUMMARY_PATH, "w") as f:
        json.dump(summary, f, indent=2)
    print("\nSUMMARY")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()