#!/usr/bin/env python3
"""HumanEval+ / MBPP+ gegen einen llama-server.

Warum diese Datensaetze: beim selbstgebauten Aufgabensatz erreichten sechs von
sieben Modellen 10/10 — ein Test, den alle bestehen, unterscheidet nichts.
Die EvalPlus-Varianten haben ein Vielfaches der Testfaelle der Originale und
sieben Loesungen aus, die nur zufaellig richtig aussehen.

Bewertet wird durch Ausfuehren, in einem eigenen Prozess mit Zeitlimit.
Kein Modell benotet ein anderes.

  ./bench_humaneval.py --port 8291 --modell "Qwen3.8 Q8_0" --satz humanevalplus \
      --out ergebnis.json
"""
import argparse
import json
import os
import re
import subprocess
import sys
import time
import urllib.request

DS = "/mnt/models/qwen38-prep/datasets"
# Die venv hat numpy, das Systempython nicht — die EvalPlus-Tests brauchen es.
PYTHON = "/mnt/models/skinnyJoe-venv/bin/python"

REASONING = False


def lade(satz, grenze):
    pfad = os.path.join(DS, f"{satz}.jsonl")
    zeilen = [json.loads(z) for z in open(pfad) if z.strip()]
    return zeilen[:grenze] if grenze else zeilen


def bau_prompt(satz, a):
    """Baut die Modellanfrage. Bewusst knapp und ohne Beispiele — sonst misst
    man die Prompt-Kunst statt das Modell."""
    if satz.startswith("humaneval"):
        return ("Vervollstaendige die folgende Python-Funktion. Gib die "
                "vollstaendige Funktion samt Signatur und noetigen Importen "
                "aus, sonst nichts:\n\n" + a["prompt"])
    # MBPP: Aufgabe in Prosa, der erwartete Funktionsname steckt im ersten
    # assert. Ohne diesen Hinweis raet das Modell den Namen und faellt durch,
    # was die Aufgabe nicht misst.
    beispiel = ""
    tl = a.get("test_list") or []
    if isinstance(tl, str):
        try:
            tl = json.loads(tl)
        except Exception:
            tl = [tl]
    if tl:
        beispiel = f"\n\nDie Funktion muss so aufrufbar sein:\n{tl[0]}"
    return (f"Schreibe eine Python-Funktion fuer diese Aufgabe:\n\n"
            f"{a['prompt']}{beispiel}\n\nGib nur den Code aus, ohne Erklaerung.")


def frage(port, prompt, max_tokens, timeout):
    b = {"messages": [{"role": "user", "content": prompt}],
         "max_tokens": max_tokens, "cache_prompt": False,
         "temperature": 0.2, "top_p": 0.9}
    if not REASONING:
        b["chat_template_kwargs"] = {"enable_thinking": False}
    req = urllib.request.Request(
        f"http://127.0.0.1:{port}/v1/chat/completions",
        data=json.dumps(b).encode(),
        headers={"Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            d = json.load(r)
    except Exception as e:
        return "", dict(tps=0.0, tokens=0, sekunden=time.time() - t0,
                        abbruch="fehler"), f"{type(e).__name__}: {e}"
    m = d["choices"][0]["message"]
    t = d.get("timings", {})
    return (m.get("content") or ""), dict(
        tps=t.get("predicted_per_second") or 0.0,
        tokens=t.get("predicted_n") or 0,
        denk_zeichen=len(m.get("reasoning_content") or ""),
        abbruch=d["choices"][0].get("finish_reason") or "",
        sekunden=round(time.time() - t0, 2)), ""


def schaele(text):
    zaun = re.findall(r"```(?:python|py)?\s*\n(.*?)```", text, re.S)
    if zaun:
        return max(zaun, key=len).strip()
    return text.strip()


def pruefe(satz, a, loesung, timeout=25):
    """Fuehrt die Loesung gegen die Testfaelle aus. Liefert (ok, grund)."""
    if not loesung.strip():
        return False, "leere Antwort"
    if satz.startswith("humaneval"):
        # Das Modell liefert die ganze Funktion. Der Prompt-Kopf kommt trotzdem
        # davor: er enthaelt die Importe, auf die der Docstring sich stuetzt.
        # Doppelte def-Zeilen stoeren nicht, die zweite gewinnt.
        programm = (a["prompt"] + "\n\npass\n\n" + loesung + "\n\n" +
                    a["test"] + f"\n\ncheck({a['entry_point']})\n")
    else:
        test = a.get("test") or ""
        vorspann = a.get("test_imports") or []
        if isinstance(vorspann, str):
            vorspann = [vorspann]
        programm = (loesung + "\n\n" + "\n".join(vorspann) + "\n" + test)
        # MBPP+ definiert die Pruefroutine, ruft sie aber nicht auf.
        if "def check(" in test and re.search(r"^check\(", test, re.M) is None:
            m = re.search(r"assert\s+\w*\(?\s*(\w+)\s*\(", test)
            programm += "\ncheck()\n" if "def check()" in test else ""
    # Ueber eine Datei statt python -c: die EvalPlus-Testbloecke sind teils
    # ueber 100 kB gross und sprengen das Argumentlimit der Kommandozeile
    # (OSError 7). Genau diese Testmenge macht sie haerter als die Originale.
    import tempfile
    fd, pfad = tempfile.mkstemp(suffix=".py", prefix="evalplus-",
                                dir="/mnt/models/tmp")
    try:
        with os.fdopen(fd, "w") as f:
            f.write(programm)
        p = subprocess.run([PYTHON, pfad], capture_output=True,
                           text=True, timeout=timeout)
    except subprocess.TimeoutExpired:
        return False, f"Zeitueberschreitung ({timeout}s)"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"
    finally:
        try:
            os.unlink(pfad)
        except OSError:
            pass
    if p.returncode == 0:
        return True, ""
    letzte = (p.stderr.strip().splitlines() or ["?"])[-1]
    return False, letzte[:100]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8291)
    ap.add_argument("--modell", required=True)
    ap.add_argument("--satz", default="humanevalplus")
    ap.add_argument("--out", required=True)
    ap.add_argument("--grenze", type=int, default=0, help="0 = alle")
    ap.add_argument("--max-tokens", type=int, default=1200)
    ap.add_argument("--timeout", type=int, default=300)
    ap.add_argument("--reasoning", action="store_true")
    args = ap.parse_args()

    global REASONING
    REASONING = args.reasoning

    aufgaben = lade(args.satz, args.grenze)
    print(f"{args.modell} | {args.satz} | {len(aufgaben)} Aufgaben | "
          f"Reasoning {'AN' if REASONING else 'AUS'}", file=sys.stderr)

    zeilen, ok_n = [], 0
    t_start = time.time()
    for i, a in enumerate(aufgaben, 1):
        txt, m, fehler = frage(args.port, bau_prompt(args.satz, a),
                               args.max_tokens, args.timeout)
        if fehler:
            zeilen.append(dict(task=str(a["task_id"]), ok=False, grund=fehler, **m))
            print(f"  [{i:3d}/{len(aufgaben)}] {a['task_id']}: SERVERFEHLER "
                  f"{fehler[:40]}", file=sys.stderr)
            continue
        loesung = schaele(txt)
        ok, grund = pruefe(args.satz, a, loesung)
        ok_n += ok
        zeilen.append(dict(task=str(a["task_id"]), ok=ok, grund=grund, **m))
        if i % 10 == 0 or not ok:
            print(f"  [{i:3d}/{len(aufgaben)}] {a['task_id']}: "
                  f"{'OK' if ok else 'FEHL ' + grund[:50]}  "
                  f"(bisher {ok_n}/{i} = {ok_n/i*100:.0f}%)", file=sys.stderr)

    dauer = time.time() - t_start
    tps = sorted(x["tps"] for x in zeilen if x["tps"] > 0)
    bericht = dict(
        modell=args.modell, satz=args.satz, reasoning=REASONING,
        n=len(aufgaben), bestanden=ok_n,
        quote=round(ok_n / len(aufgaben) * 100, 1) if aufgaben else 0.0,
        tps_median=round(tps[len(tps) // 2], 1) if tps else 0.0,
        tokens_gesamt=sum(x["tokens"] for x in zeilen),
        tokens_je_aufgabe=round(sum(x["tokens"] for x in zeilen) / len(zeilen), 1)
                          if zeilen else 0,
        am_limit=sum(1 for x in zeilen if x.get("abbruch") == "length"),
        dauer_s=round(dauer, 1),
        aufgaben=zeilen)
    with open(args.out, "w") as f:
        json.dump(bericht, f, ensure_ascii=False, indent=1)
    print(f"\n{args.modell} | {args.satz}: {ok_n}/{len(aufgaben)} = "
          f"{bericht['quote']}%  |  {bericht['tps_median']} t/s  |  "
          f"{bericht['tokens_je_aufgabe']:.0f} Token/Aufgabe  |  "
          f"{bericht['am_limit']} abgeschnitten  |  {dauer/60:.1f} min",
          file=sys.stderr)


if __name__ == "__main__":
    main()