"""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. --context FILE... puts the files' text (up to --context-chars) before every
prompt and the random tag after it: run 1 prefills the whole context, later runs find it in
the prefix cache, like an agent's system prompt and tools.
"""

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, context=""):
    text = f"{context}[{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)
    parser.add_argument("--context", nargs="*", default=[], help="files to put before each prompt")
    parser.add_argument("--context-chars", type=int, default=48000)
    args = parser.parse_args()
    model = first_model(args.base_url)
    context = ""
    for path in args.context:
        with open(path) as f:
            context += f.read()
    context = context[: args.context_chars] + "\n\n" if context else ""
    results = []
    for i in range(args.runs):
        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens, context))
        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()