#!/usr/bin/env python3
"""Faehigkeits-Benchmark: Terminal und Programmierung, gegen einen llama-server.

Bewertet wird ausschliesslich maschinell:
  * Terminal-Aufgaben — der gelieferte Shell-Einzeiler laeuft in einer
    Wegwerf-Sandbox, geprueft wird die tatsaechliche Ausgabe bzw. der
    Dateizustand danach.
  * Code-Aufgaben — die gelieferte Funktion laeuft gegen feste Testfaelle.

Kein Modell benotet ein anderes. Wo etwas nicht laeuft, gilt es als
durchgefallen, und der Grund wird protokolliert.

  ./bench_skills.py --port 8290 --modell "Qwen3.8-27B Q8_0" --out ergebnis.json
"""
import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import tasks_terminal
import tasks_code

# Wird aus main() gesetzt. Steuert, ob der Denkblock zugelassen wird.
REASONING = False

# Muster, die eine Shell-Antwort sofort disqualifizieren. Der Befehl laeuft zwar
# in einer Sandbox, aber ein 'rm -rf /' oder ein sudo hat dort trotzdem nichts
# verloren — das waere ein Treffer gegen das echte System.
GEFAEHRLICH = re.compile(
    r"(^|[;&|`(\s])(sudo|doas|su)\s"
    r"|rm\s+(-[a-zA-Z]*\s+)*(/|/\*|~|\$HOME)(\s|$)"
    r"|mkfs|fdisk|dd\s+[^|]*of=/dev/"
    r"|>\s*/dev/(sd|nvme)"
    r"|shutdown|reboot|halt|init\s+0"
    r"|:\(\)\s*\{.*\}\s*;\s*:"           # Fork-Bombe
    r"|curl[^|]*\|\s*(ba)?sh|wget[^|]*\|\s*(ba)?sh"
    r"|chmod\s+-R\s+777\s+/"
    r"|systemctl|pkill|killall"
)


def frage(port, prompt, max_tokens, timeout):
    """Eine Anfrage an den Server. Liefert (text, messwerte, fehler).

    messwerte enthaelt beide Token-Raten: die Generierrate (was der Nutzer als
    Tempo empfindet) und die Prompt-Rate (wie schnell die Eingabe verarbeitet
    wird). Nur die Generierrate zu melden verschweigt die halbe Wahrheit —
    bei langen Eingaben dominiert die Prompt-Phase die Wartezeit."""
    b = {
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "cache_prompt": False,
        "temperature": 0.2,
        "top_p": 0.9,
    }
    if not REASONING:
        # NUR ueber chat_template_kwargs. Das Feld "enable_thinking" auf
        # oberster Ebene wird von llama.cpp nicht an die Chat-Vorlage
        # durchgereicht und bleibt wirkungslos — gemessen: 825 Zeichen
        # Denkblock trotz enable_thinking=false, gegen 0 Zeichen hiermit.
        b["chat_template_kwargs"] = {"enable_thinking": False}
    body = json.dumps(b).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()
    leer = dict(tps=0.0, prompt_tps=0.0, tokens=0, prompt_tokens=0,
                sekunden=0.0, ttft=0.0)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            d = json.load(r)
    except Exception as e:
        leer["sekunden"] = time.time() - t0
        return "", leer, f"{type(e).__name__}: {e}"
    sek = time.time() - t0
    m = d.get("choices", [{}])[0].get("message", {})
    txt = m.get("content") or ""
    denken = m.get("reasoning_content") or ""
    t = d.get("timings", {})
    return txt, dict(
        tps=t.get("predicted_per_second") or 0.0,
        prompt_tps=t.get("prompt_per_second") or 0.0,
        tokens=t.get("predicted_n") or 0,
        prompt_tokens=t.get("prompt_n") or 0,
        # Zeichen statt Token, weil der Server den Denkblock nicht getrennt
        # zaehlt. Verhaeltnis Denken/Antwort ist damit trotzdem ablesbar.
        denk_zeichen=len(denken),
        antwort_zeichen=len(txt),
        abbruch=d.get("choices", [{}])[0].get("finish_reason") or "",
        # Zeit bis zum ersten Token: die Prompt-Phase in Millisekunden.
        ttft=round((t.get("prompt_ms") or 0.0) / 1000.0, 3),
        sekunden=round(sek, 2),
    ), ""


def schaele_code(text):
    """Holt den Code aus der Antwort. Modelle liefern mal mit Zaun, mal ohne,
    mal mit Vorrede — alle drei Faelle muessen durchgehen, sonst misst man
    Formatierungsdisziplin statt Koennen."""
    zaun = re.findall(r"```(?:python|py|bash|sh|shell)?\s*\n(.*?)```",
                      text, re.S)
    if zaun:
        return max(zaun, key=len).strip()
    # Kein Zaun: alles nach einer etwaigen Vorrede nehmen.
    return text.strip()


def schaele_befehl(text):
    """Holt den Shell-Einzeiler. Nimmt die erste Zeile, die nach Befehl aussieht."""
    roh = schaele_code(text)
    zeilen = [z.strip() for z in roh.splitlines() if z.strip()]
    zeilen = [z for z in zeilen if not z.startswith("#")]
    if not zeilen:
        return ""
    # Ein '$' oder '> ' am Anfang ist Prompt-Deko, kein Befehl.
    z = zeilen[0]
    z = re.sub(r"^\s*(\$|>|%)\s+", "", z)
    # Mehrzeilige Befehle mit Backslash-Fortsetzung zusammensetzen.
    if z.endswith("\\") and len(zeilen) > 1:
        z = " ".join(x.rstrip("\\").strip() for x in zeilen)
    return z.strip()


def baue_sandbox():
    d = tempfile.mkdtemp(prefix="bench-sandbox-")
    for name, inhalt in tasks_terminal.SANDBOX.items():
        with open(os.path.join(d, name), "w") as f:
            f.write(inhalt)
    for name, groesse in tasks_terminal.SANDBOX_SIZED.items():
        with open(os.path.join(d, name), "wb") as f:
            f.write(b"\0" * groesse)
    return d


def lauf_terminal(aufgabe, befehl):
    """Fuehrt den Befehl in einer frischen Sandbox aus. Liefert (bestanden, grund)."""
    if not befehl:
        return False, "leere Antwort"
    if GEFAEHRLICH.search(befehl):
        return False, "abgelehnt: gefaehrliches Muster"
    d = baue_sandbox()
    try:
        p = subprocess.run(["bash", "-c", befehl], cwd=d, timeout=15,
                           capture_output=True, text=True, errors="replace")
        out = p.stdout
        if aufgabe.get("datei_check"):
            dateien = {}
            for name in tasks_terminal.SANDBOX:
                pfad = os.path.join(d, name)
                if os.path.isfile(pfad):
                    dateien[name] = open(pfad, errors="replace").read()
            ok = bool(aufgabe["datei_check"](dateien))
            return ok, "" if ok else f"Dateizustand falsch (rc={p.returncode})"
        if p.returncode != 0 and not out.strip():
            return False, f"rc={p.returncode}: {p.stderr.strip()[:80]}"
        try:
            ok = bool(aufgabe["check"](out))
        except Exception as e:
            return False, f"Ausgabe unbrauchbar: {type(e).__name__}"
        return ok, "" if ok else f"Ausgabe falsch: {out.strip()[:60]!r}"
    except subprocess.TimeoutExpired:
        return False, "Zeitueberschreitung (15s)"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"
    finally:
        shutil.rmtree(d, ignore_errors=True)


# Wird als eigener Prozess gestartet, damit eine Endlosschleife oder ein
# Absturz im Modellcode den Benchmark nicht mitreisst.
RUNNER = r'''
import json, sys
quelle = sys.stdin.read()
auftrag = json.loads(quelle)
ns = {}
ergebnisse = []
try:
    exec(auftrag["code"], ns)
except Exception as e:
    print(json.dumps({"fehler": "Code laedt nicht: %s: %s" % (type(e).__name__, e)}))
    sys.exit(0)
for ausdruck, erwartet in auftrag["tests"]:
    try:
        lokal = dict(ns)
        # ERG als Sentinel, nicht "_r": Funktionsnamen wie suche_rotiert
        # enthalten "_r" und landeten sonst faelschlich im exec-Zweig.
        if "ERG" in ausdruck:
            exec(ausdruck, lokal)
            wert = lokal.get("ERG")
        else:
            wert = eval(ausdruck, lokal)
        # Tupel und Listen gelten als gleich, wenn der Inhalt stimmt —
        # die Aufgabe schreibt die Behaeltersorte nicht vor.
        def norm(x):
            if isinstance(x, (list, tuple)):
                return [norm(i) for i in x]
            return x
        ergebnisse.append(norm(wert) == norm(erwartet))
    except Exception as e:
        ergebnisse.append(False)
print(json.dumps({"ergebnisse": ergebnisse}))
'''


def lauf_code(aufgabe, code):
    """Laesst den gelieferten Code gegen die Testfaelle laufen."""
    if not code.strip():
        return False, 0, len(aufgabe["tests"]), "leere Antwort"
    auftrag = json.dumps({"code": code, "tests": aufgabe["tests"]})
    try:
        p = subprocess.run([sys.executable, "-c", RUNNER], input=auftrag,
                           capture_output=True, text=True, timeout=30)
    except subprocess.TimeoutExpired:
        return False, 0, len(aufgabe["tests"]), "Zeitueberschreitung (30s)"
    try:
        d = json.loads(p.stdout.strip().splitlines()[-1])
    except Exception:
        return False, 0, len(aufgabe["tests"]), f"Runner stumm: {p.stderr.strip()[:70]}"
    if "fehler" in d:
        return False, 0, len(aufgabe["tests"]), d["fehler"]
    e = d["ergebnisse"]
    return all(e), sum(e), len(e), "" if all(e) else f"{sum(e)}/{len(e)} Testfaelle"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8290)
    ap.add_argument("--modell", required=True, help="Anzeigename fuer die Tabelle")
    ap.add_argument("--out", required=True)
    ap.add_argument("--max-tokens", type=int, default=900)
    ap.add_argument("--timeout", type=int, default=300)
    ap.add_argument("--nur", default="", help="term|code, sonst beides")
    ap.add_argument("--reasoning", action="store_true",
                    help="Denkblock zulassen (sonst per chat_template_kwargs aus)")
    args = ap.parse_args()

    global REASONING
    REASONING = args.reasoning
    print(f"Reasoning: {'AN' if REASONING else 'AUS'}, "
          f"max_tokens={args.max_tokens}", file=sys.stderr)

    bericht = {"modell": args.modell, "reasoning": REASONING,
               "max_tokens": args.max_tokens,
               "terminal": [], "code": [], "messungen": []}

    if args.nur in ("", "term"):
        print(f"--- Terminal ({len(tasks_terminal.TASKS)} Aufgaben) ---",
              file=sys.stderr)
        for a in tasks_terminal.TASKS:
            txt, m, fehler = frage(args.port, a["prompt"],
                                   args.max_tokens, args.timeout)
            if fehler:
                bericht["terminal"].append(
                    dict(id=a["id"], titel=a["titel"], ok=False, grund=fehler,
                         antwort=""))
                print(f"  {a['id']} {a['titel']:<26} SERVERFEHLER {fehler[:50]}",
                      file=sys.stderr)
                continue
            befehl = schaele_befehl(txt)
            ok, grund = lauf_terminal(a, befehl)
            bericht["terminal"].append(
                dict(id=a["id"], titel=a["titel"], ok=ok, grund=grund,
                     antwort=befehl[:200], **m))
            bericht["messungen"].append(m)
            print(f"  {a['id']} {a['titel']:<26} {'OK ' if ok else 'FEHL'} "
                  f"{m['tps']:5.1f} t/s  {m['tokens']:4d} Tok  {grund[:45]}",
                  file=sys.stderr)

    if args.nur in ("", "code"):
        print(f"--- Programmierung ({len(tasks_code.TASKS)} Aufgaben) ---",
              file=sys.stderr)
        for a in tasks_code.TASKS:
            txt, m, fehler = frage(args.port, a["prompt"],
                                   args.max_tokens, args.timeout)
            if fehler:
                bericht["code"].append(
                    dict(id=a["id"], titel=a["titel"], ok=False, grund=fehler,
                         bestanden=0, gesamt=len(a["tests"])))
                print(f"  {a['id']} {a['titel']:<26} SERVERFEHLER {fehler[:50]}",
                      file=sys.stderr)
                continue
            code = schaele_code(txt)
            ok, best, ges, grund = lauf_code(a, code)
            bericht["code"].append(
                dict(id=a["id"], titel=a["titel"], ok=ok, grund=grund,
                     bestanden=best, gesamt=ges, code=code[:600], **m))
            bericht["messungen"].append(m)
            print(f"  {a['id']} {a['titel']:<26} {'OK ' if ok else 'FEHL'} "
                  f"{m['tps']:5.1f} t/s  {m['tokens']:4d} Tok  "
                  f"{best}/{ges} {grund[:35]}", file=sys.stderr)

    t, c, mm = bericht["terminal"], bericht["code"], bericht["messungen"]
    tps = sorted(x["tps"] for x in mm if x["tps"] > 0)
    ptps = sorted(x["prompt_tps"] for x in mm if x["prompt_tps"] > 0)
    med = lambda v: v[len(v) // 2] if v else 0.0
    bericht["summe"] = {
        "terminal_ok": sum(1 for x in t if x["ok"]), "terminal_n": len(t),
        "code_ok": sum(1 for x in c if x["ok"]), "code_n": len(c),
        "code_testfaelle_ok": sum(x.get("bestanden", 0) for x in c),
        "code_testfaelle_n": sum(x.get("gesamt", 0) for x in c),
        "tps_median": round(med(tps), 1),
        "tps_min": round(tps[0], 1) if tps else 0.0,
        "tps_max": round(tps[-1], 1) if tps else 0.0,
        "prompt_tps_median": round(med(ptps), 1),
        "tokens_gesamt": sum(x["tokens"] for x in mm),
        "sekunden_gesamt": round(sum(x["sekunden"] for x in mm), 1),
        "anfragen": len(mm),
        # Der Preis des Denkens: wie viele Token im Schnitt, und wie viel davon
        # Grubelei statt Antwort war. Ein Lauf, der ans Token-Limit stoesst,
        # hat die Antwort nie ausgegeben — deshalb wird das mitgezaehlt.
        "tokens_je_aufgabe": round(sum(x["tokens"] for x in mm) / len(mm), 1)
                             if mm else 0,
        "denk_zeichen_gesamt": sum(x.get("denk_zeichen", 0) for x in mm),
        "antwort_zeichen_gesamt": sum(x.get("antwort_zeichen", 0) for x in mm),
        "am_limit": sum(1 for x in mm if x.get("abbruch") == "length"),
    }
    with open(args.out, "w") as f:
        json.dump(bericht, f, ensure_ascii=False, indent=1)
    s = bericht["summe"]
    print(f"\n{args.modell} [Reasoning {'AN' if REASONING else 'AUS'}]: "
          f"Terminal {s['terminal_ok']}/{s['terminal_n']}, "
          f"Code {s['code_ok']}/{s['code_n']} "
          f"({s['code_testfaelle_ok']}/{s['code_testfaelle_n']} Testfaelle), "
          f"{s['tps_median']:.1f} t/s Median "
          f"({s['tps_min']:.1f}-{s['tps_max']:.1f}), "
          f"Prompt {s['prompt_tps_median']:.0f} t/s, "
          f"{s['tokens_gesamt']} Token in {s['sekunden_gesamt']}s, "
          f"{s['tokens_je_aufgabe']:.0f} Token/Aufgabe, "
          f"{s['am_limit']} am Limit abgeschnitten", file=sys.stderr)
    print(f"geschrieben: {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()