#!/usr/bin/env python3
"""Misst die Prompt-Verarbeitungsrate mit einem langen Prompt.

Warum die Prompt-Rate und nicht die Token-Rate: der stille CPU-Fallback bei
quantisiertem KV-Cache trifft die Attention-Berechnung. Die schlaegt beim
Einlesen des Prompts voll durch (alle Token auf einmal), bei der Erzeugung
dagegen kaum (ein Token nach dem anderen, speicherbandbreiten-limitiert).
Wer nur t/s misst, uebersieht den Fehler.
"""
import json
import sys
import time
import urllib.request

PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8291
NAME = sys.argv[2] if len(sys.argv) > 2 else "?"

# Rund 6000 Token Fliesstext — lang genug, dass die Prompt-Phase dominiert.
BAUSTEIN = (
    "Der Rhein ist mit 1233 Kilometern einer der laengsten Fluesse Europas und "
    "verbindet die Alpen mit der Nordsee. Fuer die Industrialisierung "
    "Deutschlands war er die entscheidende Verkehrsader: Kohle aus dem "
    "Ruhrgebiet, Erz aus Lothringen und Getreide aus dem Oberrheingraben "
    "wurden auf ihm bewegt, lange bevor die Eisenbahn diese Mengen bewaeltigen "
    "konnte. Die Staedte entlang des Flusses wuchsen entsprechend schnell. "
)
PROMPT = BAUSTEIN * 90 + "\n\nFasse den Text in genau einem Satz zusammen."


def einmal(max_tokens=40):
    body = json.dumps({
        "messages": [{"role": "user", "content": PROMPT}],
        "max_tokens": max_tokens,
        "cache_prompt": False,
        "temperature": 0.2,
        "chat_template_kwargs": {"enable_thinking": False},
    }).encode()
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}/v1/chat/completions",
        data=body, headers={"Content-Type": "application/json"})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=900) as r:
        d = json.load(r)
    wand = time.time() - t0
    t = d.get("timings", {})
    return (t.get("prompt_n") or 0, t.get("prompt_per_second") or 0.0,
            t.get("predicted_per_second") or 0.0, wand)


try:
    # Erster Lauf waermt die Puffer, gewertet wird der zweite.
    einmal(8)
    n, ptps, tps, wand = einmal()
    print(f"{NAME:14s}  Prompt {n:6d} Tok  "
          f"{ptps:9.1f} Tok/s prefill  |  {tps:5.1f} Tok/s generate  "
          f"|  {wand:6.1f}s")
except Exception as e:
    print(f"{NAME:14s}  FEHLER {type(e).__name__}: {str(e)[:70]}")