"""Time to first token and decode speed of an OpenAI-compatible server, as JSON lines.

usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3]
Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly
--tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds
the medians.
"""

import argparse
import json
import statistics
import time
import uuid

from prain.chat import chat_body, first_model, post, stream_text

PROMPTS = [
    "Explain how a CPU cache works, step by step, for a beginner.",
    "Write a short story about a lighthouse keeper who finds a message in a bottle.",
    "Schreibe eine kurze Anleitung, wie man Brot backt.",
]


def run(base_url, model, prompt, tokens):
    text = f"[{uuid.uuid4().hex[:8]}] {prompt}"
    body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens,
                     temperature=0, ignore_eos=True)
    start = time.monotonic()
    stamps = []
    with post(base_url, "/v1/chat/completions", body, timeout=3600) as response:
        for _ in stream_text(response):
            stamps.append(time.monotonic())
    return {"ttft_s": round(stamps[0] - start, 3),
            "decode_tok_s": round((tokens - 1) / max(stamps[-1] - stamps[0], 1e-9), 3)}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("base_url")
    parser.add_argument("--tokens", type=int, default=128)
    parser.add_argument("--runs", type=int, default=3)
    args = parser.parse_args()
    model = first_model(args.base_url)
    results = []
    for i in range(args.runs):
        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens))
        print(json.dumps(results[-1]), flush=True)
    print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]}
                     | {"median_of": len(results), "tokens": args.tokens}))


if __name__ == "__main__":
    main()