"""Benchmark: Referenzaufgaben mit bekanntem Optimum durch die Pipeline schicken und Modelle vergleichen.

Beispiel:
    .venv/bin/python benchmark.py --models deepseek/deepseek-v4-flash-0731 deepseek/deepseek-v4-pro-0813 --runs 2
"""
import argparse
import hashlib
import json
import os
import subprocess
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
from datetime import datetime
from pathlib import Path

from core.llm import DEFAULT_MODEL, is_fatal_error
from TextToModel import run_pipeline

BENCHMARK_DIR = Path(__file__).parent
ROOT = BENCHMARK_DIR.parent
CASES_DIR = BENCHMARK_DIR / "cases"
RESULTS_DIR = BENCHMARK_DIR / "results"

# Nach dem ersten fatalen API-Fehler (Kontingent/Guthaben/Key) keine weiteren Läufe starten
quota_exhausted = threading.Event()

# Dateien, die das LLM-Verhalten steuern: die System-Prompts und die Feldbeschreibungen
# (Field(description=...) geht als JSON-Schema mit und wirkt wie eine Prompt-Anweisung).
PROMPT_FILES = ("understanding/Verstehen-Prompt.txt", "understanding/Abbildungs-Prompt.txt",
                "modeling/Routing-und-Extraktions-Prompt.txt", "checks/Pruef-Prompt.txt",
                "solving/Solver-Prompt.txt",
                "core/model_schema.py")

# Umgebungsvariablen, die das Verhalten ändern – ohne sie ließe sich ein A/B-Lauf nicht zuordnen
CONFIG_VARS = ("LLM_MODEL", "LLM_MODEL_EXTRACTION", "LLM_MODEL_REVIEW", "LLM_MODEL_REVISION", "LLM_MODEL_SOLVER",
               "UNDERSTAND_RETRIES", "MAX_REVISIONS", "SOLVER_RETRIES", "MAX_RESTARTS", "INFEASIBLE_RECHECK",
               "LLM_REASONING_EFFORT", "SOLVER_SANDBOX")


def prompt_digest(paths=None) -> str:
    """Kurzer Hash über alle Dateien, die das LLM-Verhalten steuern."""
    digest = hashlib.sha256()
    for path in paths if paths is not None else [ROOT / n for n in PROMPT_FILES]:
        digest.update(Path(path).read_bytes())
    return digest.hexdigest()[:12]


@lru_cache(maxsize=1)
def prompt_state() -> dict:
    """Womit wurde gemessen? Ohne diese Angaben lassen sich zwei Ergebnisdateien nicht vergleichen.

    Beobachtet: die Messlatte 19/22 ließ sich keinem Prompt-Stand zuordnen, weil nur Fall, Modell und
    Laufnummer gespeichert wurden – jede Prompt-Änderung entwertete die Vergleichbarkeit still.
    """
    state = {"prompts": prompt_digest()}
    gesetzt = {name: os.environ[name] for name in CONFIG_VARS if os.environ.get(name)}
    if gesetzt:
        state["config"] = gesetzt
    def git(*args):
        return subprocess.run(["git", *args], cwd=ROOT, capture_output=True,
                              text=True, timeout=10).stdout.strip()
    try:
        state["commit"] = git("rev-parse", "--short", "HEAD") or "unbekannt"
        state["dirty"] = bool(git("status", "--porcelain"))
    except (OSError, subprocess.SubprocessError):
        state["commit"], state["dirty"] = "unbekannt", None
    return state


def load_cases(names=None):
    cases = []
    for case_dir in sorted(p for p in CASES_DIR.iterdir() if p.is_dir()):
        if names and case_dir.name not in names:
            continue
        cases.append({
            "name": case_dir.name,
            "task": (case_dir / "task.txt").read_text(encoding="utf-8"),
            "expected": json.loads((case_dir / "expected.json").read_text(encoding="utf-8")),
            # Alles außer task.txt/expected.json gilt als beigelegte Datendatei (--daten)
            "data_files": sorted(str(f) for f in case_dir.iterdir()
                                 if f.is_file() and f.name not in ("task.txt", "expected.json")),
        })
    return cases


def result_correct(solver_run, expected: dict) -> bool:
    """Status muss passen (Default optimal); bei optimal zusätzlich der Zielwert."""
    expected_status = expected.get("status", "optimal")
    if solver_run is None or solver_run.status.value != expected_status:
        return False
    return expected_status != "optimal" or objective_matches(solver_run.objective_value, expected)


def objective_matches(actual, expected: dict) -> bool:
    if actual is None:
        return False
    # Toleranz relativ, mindestens absolut (für Zielwerte < 1)
    return abs(actual - expected["objective_value"]) <= expected["tolerance"] * max(1.0, abs(expected["objective_value"]))


def run_case(case: dict, model: str, run_no: int) -> dict:
    expected = case["expected"]
    row = {"case": case["name"], "model": model, "run": run_no, **prompt_state()}
    if quota_exhausted.is_set():
        row.update(seconds=0, error="übersprungen: fataler API-Fehler in früherem Lauf")
        return row
    start = time.monotonic()
    try:
        result = run_pipeline(case["task"], model=model, data_files=case.get("data_files") or ())
    except Exception as exc:  # API-/Validierungsfehler sollen den Benchmark nicht abbrechen
        message = f"{type(exc).__name__}: {exc}"
        if is_fatal_error(message):
            quota_exhausted.set()
        row.update(seconds=round(time.monotonic() - start), error=message[:500])
        return row

    spec, solver_run = result.specification, result.solver_run
    if spec is None:  # abgelehnt: keine Optimierungsaufgabe erkannt
        fragen = len(result.understanding.sub_questions) if result.understanding else 0
        row.update(seconds=round(time.monotonic() - start), problem_class=None, target_stack=None,
                   class_ok=False, stack_ok=False, check_errors=0, check_warnings=len(result.checks),
                   revisions=0, review_ok=None, data_assumptions=0, solver_status="abgelehnt", solver_attempts=0,
                   objective_value=None, objective_ok=False, questions=fragen, answered=0, answer_findings=0,
                   task_ok=False, solver_error="; ".join(result.reliability_issues), reliable=False, restarts=0,
                   providers=result.usage.providers if result.usage else {},
                   responded_models=result.usage.responded_models if result.usage else {},
                   llm_calls=result.usage.calls if result.usage else 0,
                   reasoning_tokens=result.usage.reasoning_tokens if result.usage else 0,
                   cost_usd=result.usage.cost_usd if result.usage else None,
                   result=json.loads(result.model_dump_json()))
        return row
    # "Aufgabe erfüllt" statt nur "Zahl richtig": alle Teilfragen beantwortet, keine unbelegte Zahl
    fragen = {q.id for q in result.understanding.sub_questions} if result.understanding else set()
    antworten = {a.question_id for a in solver_run.answers} if solver_run else set()
    abnahme = [w for w in (solver_run.warnings if solver_run else []) if "Antwort zu" in w or "Teilfragen ohne" in w]
    objective_ok = result.reliable and result_correct(solver_run, expected)
    row.update(
        seconds=round(time.monotonic() - start),
        problem_class=spec.problem_class.value,
        target_stack=spec.target_stack.value,
        class_ok=spec.problem_class.value == expected.get("problem_class", spec.problem_class.value),
        stack_ok=spec.target_stack.value == expected.get("target_stack", spec.target_stack.value),
        check_errors=sum(1 for i in result.checks if i.severity.value == "error"),
        check_warnings=sum(1 for i in result.checks if i.severity.value == "warning"),
        revisions=result.revisions,
        review_ok=not result.review.needs_revision,
        data_assumptions=len(result.solver_code.data_assumptions) if result.solver_code else 0,
        solver_status=solver_run.status.value if solver_run else "nicht gestartet",
        solver_attempts=solver_run.attempts if solver_run else 0,
        objective_value=solver_run.objective_value if solver_run else None,
        objective_ok=objective_ok,
        questions=len(fragen),
        answered=len(fragen & antworten),
        answer_findings=len(abnahme),
        task_ok=bool(objective_ok and fragen and fragen <= antworten and not abnahme
                     and len(fragen) >= expected.get("min_sub_questions", 1)),
        solver_error=solver_run.error if solver_run else "; ".join(result.reliability_issues),
        reliable=result.reliable,
        restarts=result.restarts,
        providers=result.usage.providers,
        responded_models=result.usage.responded_models,
        llm_calls=result.usage.calls,
        reasoning_tokens=result.usage.reasoning_tokens,
        cost_usd=result.usage.cost_usd,
        result=json.loads(result.model_dump_json()),
    )
    return row


def mark(value) -> str:
    return {True: "ja", False: "NEIN", None: "-"}.get(value, str(value))


def print_table(rows):
    header = ["Fall", "Modell", "#", "Klasse", "Stack", "ChkErr", "Rev", "Review", "Solver", "Zielwert", "Zuverl.", "Neust.", "Treffer", "Fragen", "Antw.", "Erfüllt", "Datenann.", "Aufrufe", "Reasoning", "USD", "s"]
    lines = [header]
    for r in rows:
        if "error" in r:
            lines.append([r["case"], r["model"], str(r["run"]), "-", "-", "-", "-", "-", "ABBRUCH", "-", "-", "-", "NEIN", "-", "-", "NEIN", "-", "-", "-", "-", str(r["seconds"])])
            continue
        objective = "-" if r["objective_value"] is None else f"{r['objective_value']:.6g}"
        lines.append([
            r["case"], r["model"], str(r["run"]), mark(r["class_ok"]), mark(r["stack_ok"]), str(r["check_errors"]),
            str(r["revisions"]), mark(r["review_ok"]), r["solver_status"], objective, mark(r["reliable"]), str(r["restarts"]), mark(r["objective_ok"]),
            str(r.get("questions", 0)), str(r.get("answered", 0)), mark(r.get("task_ok")),
            str(r["data_assumptions"]), str(r["llm_calls"]), str(r["reasoning_tokens"]),
            "-" if r["cost_usd"] is None else f"{r['cost_usd']:.4f}", str(r["seconds"]),
        ])
    widths = [max(len(line[i]) for line in lines) for i in range(len(header))]
    for n, line in enumerate(lines):
        print("  ".join(cell.ljust(w) for cell, w in zip(line, widths)))
        if n == 0:
            print("  ".join("-" * w for w in widths))

    stand = prompt_state()
    geaendert = " + nicht eingecheckte Änderungen" if stand.get("dirty") else ""
    print(f"\nGemessener Stand: Commit {stand['commit']}{geaendert}, Prompt-Kennung {stand['prompts']} "
          f"(über {', '.join(PROMPT_FILES)})")
    if stand.get("config"):
        print("Abweichende Einstellungen: " + ", ".join(f"{k}={v}" for k, v in stand["config"].items()))
    print("\nZusammenfassung je Modell:")
    for model in dict.fromkeys(r["model"] for r in rows):
        mine = [r for r in rows if r["model"] == model]
        hits = sum(1 for r in mine if r.get("objective_ok"))
        erfuellt = sum(1 for r in mine if r.get("task_ok"))
        routing = sum(1 for r in mine if r.get("class_ok") and r.get("stack_ok"))
        seconds = sum(r["seconds"] for r in mine) / len(mine)
        cost = sum(r.get("cost_usd") or 0 for r in mine)
        print(f"  {model}: Ergebnis korrekt {hits}/{len(mine)}, Aufgabe erfüllt {erfuellt}/{len(mine)}, Routing korrekt {routing}/{len(mine)}, "
              f"Ø {seconds:.0f} s pro Lauf, Kosten gesamt ${cost:.4f} (ohne abgebrochene Läufe)")
        antwortend = {}
        for r in mine:
            for name, n in (r.get("responded_models") or {}).items():
                antwortend[name] = antwortend.get(name, 0) + n
        if antwortend and set(antwortend) != {model}:
            print(f"    Modell laut Antwort: {antwortend}")

    for r in rows:
        if r.get("error") or r.get("solver_error"):
            print(f"\n[{r['case']} | {r['model']} | #{r['run']}] {r.get('error') or r.get('solver_error')}")
    if quota_exhausted.is_set():
        print("\nABBRUCH: fataler API-Fehler (Kontingent, Guthaben oder Key) – restliche Läufe wurden übersprungen.")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--models", nargs="+", default=[DEFAULT_MODEL])
    parser.add_argument("--runs", type=int, default=1, help="Läufe je Fall und Modell (Streuung sichtbar machen)")
    parser.add_argument("--cases", nargs="+", help="nur diese Fälle (Verzeichnisnamen)")
    parser.add_argument("--jobs", type=int, default=1, help="parallele Läufe (Gemini Free Tier: 1)")
    args = parser.parse_args()

    cases = load_cases(args.cases)
    jobs = [(case, model, run_no) for model in args.models for case in cases for run_no in range(1, args.runs + 1)]
    with ThreadPoolExecutor(max_workers=args.jobs) as pool:
        rows = list(pool.map(lambda job: run_case(*job), jobs))

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    out_file = RESULTS_DIR / f"{datetime.now():%Y%m%d-%H%M%S}.json"
    out_file.write_text(json.dumps(rows, indent=2, ensure_ascii=False), encoding="utf-8")

    print()
    print_table(rows)
    print(f"\nDetails: {out_file}")


if __name__ == "__main__":
    main()