#!/usr/bin/env python3
"""MTP-Messreihe gegen einen llama-server. Deutsche Prompts, nach Textsorte
getrennt, weil die Draft-Trefferquote genau daran haengt.

Der Server meldet pro Anfrage draft_n / draft_n_accepted in `timings` — daraus
kommen die Zahlen, nichts wird geschaetzt. Jeder Prompt laeuft mehrfach, ausgewertet
wird der Median, weil Einzelmessungen um +/-6 Prozent schwanken.

  ./bench_mtp.py --port 8270 --runs 3 --out ergebnis.md
"""
import argparse
import json
import statistics
import sys
import time
import urllib.request

# Textsorten, bei denen MTP erfahrungsgemaess sehr unterschiedlich trifft.
PROMPTS = [
    ("Aufzaehlung",
     "Zaehle von 1 bis 60, nur die Zahlen mit Komma getrennt, sonst nichts."),
    ("Struktur/JSON",
     "Gib ein JSON-Objekt mit 12 deutschen Staedten aus, je mit den Feldern "
     "name, bundesland, einwohner. Nur das JSON, kein Text drumherum."),
    ("Code",
     "Schreibe eine Python-Funktion, die eine CSV-Datei einliest, nach der "
     "dritten Spalte sortiert und als JSON zurueckgibt. Nur Code."),
    ("Fliesstext",
     "Schreibe drei Absaetze darueber, warum der Rhein fuer die deutsche "
     "Industriegeschichte wichtig war."),
    ("Fachtext",
     "Erklaere in etwa 300 Woertern den Unterschied zwischen einem "
     "Asynchronmotor und einem Synchronmotor."),
    ("Kreativ",
     "Schreibe eine kurze Geschichte ueber einen Leuchtturmwaerter, der eine "
     "Flaschenpost findet."),
    ("Uebersetzung",
     "Uebersetze ins Deutsche: 'The quick assessment of the situation led the "
     "engineers to redesign the cooling system, which had failed twice during "
     "the summer of the previous year.'"),
    ("Wiederholung",
     "Wiederhole den Satz 'Die Maschine laeuft stabil.' genau 25 Mal, "
     "jeweils in einer neuen Zeile."),
]


def ask(port, prompt, max_tokens, timeout):
    body = json.dumps({
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "enable_thinking": False,
        "cache_prompt": False,
        "temperature": 0.3,
    }).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=timeout) as r:
        d = json.load(r)
    wall = time.time() - t0
    t = d.get("timings", {})
    return {
        "tps": t.get("predicted_per_second"),
        "n": t.get("predicted_n"),
        "draft_n": t.get("draft_n") or 0,
        "draft_acc": t.get("draft_n_accepted") or 0,
        "wall": wall,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8270)
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--max-tokens", type=int, default=400)
    ap.add_argument("--timeout", type=int, default=600)
    ap.add_argument("--label", default="")
    ap.add_argument("--out", default="")
    args = ap.parse_args()

    rows = []
    for name, prompt in PROMPTS:
        tps, accs, drafted, accepted = [], [], 0, 0
        for i in range(args.runs):
            try:
                r = ask(args.port, prompt, args.max_tokens, args.timeout)
            except Exception as e:
                print(f"  {name} Lauf {i+1}: FEHLER {e}", file=sys.stderr)
                continue
            if r["tps"]:
                tps.append(r["tps"])
            drafted += r["draft_n"]
            accepted += r["draft_acc"]
            if r["draft_n"]:
                accs.append(r["draft_acc"] / r["draft_n"])
            print(f"  {name} Lauf {i+1}: {r['tps']:.1f} t/s, "
                  f"{r['draft_acc']}/{r['draft_n']} Draft-Treffer", file=sys.stderr)
        if not tps:
            continue
        rows.append({
            "sorte": name,
            "tps": statistics.median(tps),
            "tps_min": min(tps),
            "tps_max": max(tps),
            "quote": (accepted / drafted) if drafted else 0.0,
            "drafted": drafted,
            "accepted": accepted,
        })

    hdr = f"### Messreihe {args.label}\n\n" if args.label else "### Messreihe\n\n"
    md = hdr
    md += "| Textsorte | t/s (Median) | Spanne | Draft-Trefferquote | angenommen / erzeugt |\n"
    md += "|---|---|---|---|---|\n"
    for r in rows:
        md += (f"| {r['sorte']} | **{r['tps']:.1f}** | {r['tps_min']:.1f}–{r['tps_max']:.1f} "
               f"| {r['quote']*100:.1f} % | {r['accepted']} / {r['drafted']} |\n")
    if rows:
        gd = sum(r["drafted"] for r in rows)
        ga = sum(r["accepted"] for r in rows)
        if gd:
            md += (f"\nGesamt: **{ga/gd*100:.1f} %** Trefferquote "
                   f"({ga} / {gd} Draft-Token), {args.runs} Laeufe je Textsorte, "
                   f"max_tokens={args.max_tokens}.\n")
        else:
            # Lauf ohne Draft-Modell: es gibt nichts zu treffen.
            tps_all = [r["tps"] for r in rows]
            md += (f"\nOhne Spekulation — keine Draft-Token. "
                   f"Median ueber alle Textsorten: **{statistics.median(tps_all):.1f} t/s** "
                   f"(Spanne {min(tps_all):.1f}–{max(tps_all):.1f}), "
                   f"{args.runs} Laeufe je Textsorte, max_tokens={args.max_tokens}.\n")

    print(md)
    if args.out:
        open(args.out, "w").write(md)
        print(f"geschrieben: {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()