"""
llmpm benchmark <model> --tasks task1,task2 ...

Run standard LLM evaluation benchmarks against an installed model.
Uses EleutherAI's lm-evaluation-harness (lm_eval) under the hood.

Examples:
  llmpm benchmark meta-llama/Llama-3.2-3B-Instruct --tasks ifeval,hellaswag
  llmpm benchmark pretrained=gpt2 --tasks mmlu --num-fewshot 5
  llmpm benchmark meta-llama/Llama-3.2-3B-Instruct --tasks openllm
  llmpm benchmark unsloth/Llama-3.2-3B-Instruct-GGUF --tasks hellaswag --limit 100
  llmpm benchmark --list-tasks
  # lm_eval CLI used: lm_eval --model hf --model_args pretrained=... --tasks ...
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import tempfile
import time
from datetime import datetime
from pathlib import Path

# Pattern that matches lm-eval's "please install via pip install lm-eval[X]" messages.
_MISSING_EXTRA_RE = re.compile(r"pip install lm[-_]eval\[([^\]]+)\]", re.IGNORECASE)

import click
import requests

from llmpm import display
from llmpm.commands._picker import pick_installed_model
from llmpm.core import registry

# ── gated task substitutions ──────────────────────────────────────────────────
# Tasks that pull gated HuggingFace datasets are auto-replaced with open equivalents.
_GATED_SUBSTITUTIONS: dict[str, str] = {
    # leaderboard_gpqa inside the official "leaderboard" group uses Idavidrein/gpqa (gated).
    # leaderboard_open is the bundled replacement with all-open datasets.
    "leaderboard": "leaderboard_open",
}

# ── bundled custom tasks (shipped with llmpm) ─────────────────────────────────
# Task YAML files live in llmpm/tasks/<name>/<name>.yaml and are passed to
# lm_eval via --include_path so they work without being in lm_eval's own tree.
_BUNDLED_TASKS_DIR = Path(__file__).parent.parent / "tasks"
_BUNDLED_TASK_NAMES: frozenset[str] = frozenset(
    p.name for p in _BUNDLED_TASKS_DIR.iterdir() if p.is_dir()
) if _BUNDLED_TASKS_DIR.exists() else frozenset()

# ── well-known benchmarks (task name → short description) ─────────────────────
# Any task name from lm_eval/tasks/ is accepted; these are the most common ones.

KNOWN_TASKS: dict[str, str] = {
    # ── Leaderboard suites (most commonly used shortcuts) ─────────────────────
    "openllm":            "Open LLM Leaderboard v1 (ARC · HellaSwag · TruthfulQA · MMLU · Winogrande · GSM8K)",
    "leaderboard":        "Open LLM Leaderboard v2 — IFEval · BBH · MATH Hard · MuSR · MMLU-Pro  [no gated data]",
    "tinyBenchmarks":     "tinyBenchmarks — efficient subset proxies for 6 standard benchmarks",
    "metabench":          "Metabench — meta-benchmark across multiple eval suites",
    # ── Core benchmarks (appear in nearly every eval paper) ───────────────────
    "hellaswag":          "HellaSwag — sentence-completion commonsense (10-shot)",
    "gpqa":               "GPQA Diamond — 198 expert-level science QA (biology, chemistry, physics) [open mirror]",
    "hle":                "Humanity's Last Exam — 3,000 expert-level questions across science, math & humanities",
    "mmlu":               "MMLU — 57-subject academic knowledge (5-shot)",
    "arc_challenge":      "ARC Challenge — 4-choice science QA, hard set (25-shot)",
    "arc_easy":           "ARC Easy — 4-choice science QA, easy set (0-shot)",
    "gsm8k":              "GSM8K — grade-school math word problems (5-shot)",
    "truthfulqa":         "TruthfulQA — truthfulness & factuality (0-shot)",
    "winogrande":         "Winogrande — Winograd-schema commonsense (5-shot)",
    "ifeval":             "IFEval — verifiable instruction-following evaluation (0-shot)",
    "bbh":                "BIG-Bench Hard — 23 hard reasoning tasks (3-shot CoT)",
    "mmlu_pro":           "MMLU-Pro — harder 10-option MMLU variant (5-shot)",
    # ── Math & reasoning ─────────────────────────────────────────────────────
    "hendrycks_math":     "MATH — competition-level math problems (4-shot)",
    "minerva_math":       "Minerva Math — math reasoning with chain-of-thought (4-shot)",
    "gsm_plus":           "GSM+ — augmented GSM8K with harder variants",
    "mathqa":             "MathQA — algebraic & scientific word problems (0-shot)",
    "asdiv":              "ASDiv — diverse arithmetic word problems",
    "arithmetic":         "Arithmetic — basic n-digit arithmetic operations (0-shot)",
    "bigbench":           "BIG-Bench — 200+ diverse tasks covering reasoning, knowledge, and language",
    "anli":               "Adversarial NLI — robust NLI across 3 rounds (0-shot)",
    # ── Code generation ───────────────────────────────────────────────────────
    "humaneval":          "HumanEval — 164 Python programming problems (pass@k)",
    "mbpp":               "MBPP — 374 crowd-sourced Python programming problems",
    # ── Reading comprehension ─────────────────────────────────────────────────
    "drop":               "DROP — discrete reasoning over paragraphs (3-shot)",
    "squadv2":            "SQuAD v2 — reading comprehension with unanswerable Qs (1-shot)",
    "race":               "RACE — reading comprehension (0-shot)",
    "coqa":               "CoQA — conversational QA (0-shot)",
    "super_glue":         "SuperGLUE — 8-task language understanding suite",
    "glue":               "GLUE — 9-task NLU benchmark (validation sets)",
    # ── Knowledge & factuality ────────────────────────────────────────────────
    "triviaqa":           "TriviaQA — trivia open-domain QA (0-shot)",
    "nq_open":            "Natural Questions — open-domain QA (0-shot)",
    "sciq":               "SciQ — science exam QA with supporting documents (0-shot)",
    "webqs":              "WebQuestions — Freebase-grounded QA (0-shot)",
    "agieval":            "AGIEval — human-centric standardized exam tasks",
    # ── Commonsense reasoning ─────────────────────────────────────────────────
    "piqa":               "PIQA — physical intuition QA (0-shot)",
    "siqa":               "SocialIQA — social-interaction reasoning (0-shot)",
    "openbookqa":         "OpenBookQA — open-book science QA (0-shot)",
    "commonsense_qa":     "CommonsenseQA — commonsense multiple-choice (0-shot)",
    "wsc273":             "WSC273 — Winograd Schema Challenge (0-shot)",
    "logiqa":             "LogiQA — logical reading comprehension (0-shot)",
    "logiqa2":            "LogiQA 2.0 — updated logical reading comprehension (0-shot)",
    "swag":               "SWAG — grounded commonsense inference (0-shot)",
    "babi":               "bAbI — 20 QA tasks for reasoning over short stories",
    # ── Long context ──────────────────────────────────────────────────────────
    "longbench":          "LongBench — bilingual long-context understanding benchmark",
    # ── Language modeling & perplexity ────────────────────────────────────────
    "wikitext":           "WikiText — word-level perplexity on WikiText-103",
    "lambada":            "LAMBADA — last-word prediction requiring passage understanding",
    # ── Medical & scientific ──────────────────────────────────────────────────
    "medqa":              "MedQA — US medical licensing exam (USMLE) QA",
    "medmcqa":            "MedMCQA — Indian medical entrance exam QA",
    "pubmedqa":           "PubMedQA — biomedical research question answering",
    # ── Safety & bias ─────────────────────────────────────────────────────────
    "wmdp":               "WMDP — hazardous knowledge benchmark (bio, chem, cyber)",
    "bbq":                "BBQ — bias benchmark for QA across 9 social categories",
    "crows_pairs":        "CrowS-Pairs — measuring social biases in language models",
    "hendrycks_ethics":   "ETHICS — moral knowledge across 5 ethical frameworks",
    "realtoxicityprompts": "RealToxicityPrompts — toxicity of model continuations",
    # ── Multilingual ──────────────────────────────────────────────────────────
    "mgsm":               "MGSM — multilingual grade-school math (11 languages)",
    "belebele":           "Belebele — multilingual reading comprehension (122 languages)",
    "global_mmlu":        "Global MMLU — MMLU translated into 42 languages",
    "xnli":               "XNLI — cross-lingual NLI in 15 languages",
    "xcopa":              "XCOPA — cross-lingual commonsense (11 languages)",
    "xwinograd":          "XWinograd — multilingual Winograd schemas",
    "okapi":              "Okapi — multilingual instruction-following benchmark",
    "ceval":              "C-Eval — Chinese academic knowledge benchmark",
    "cmmlu":              "CMMLU — Chinese multi-subject language understanding",
    "kmmlu":              "KMMLU — Korean multi-subject language understanding",
    # ── Summarization ─────────────────────────────────────────────────────────
    "cnn_dailymail":      "CNN/DailyMail — news article summarization",
    
}


def _read_model_max_length(model_path: str) -> int | None:
    """Read max_position_embeddings (or n_positions) from a local config.json."""
    cfg = Path(model_path) / "config.json"
    if not cfg.exists():
        return None
    try:
        with open(cfg, encoding="utf-8") as fh:
            data = json.load(fh)
        # Try common keys used by different architectures
        for key in ("max_position_embeddings", "n_positions", "max_seq_len", "seq_length"):
            val = data.get(key)
            if isinstance(val, int) and val > 0:
                return val
    except Exception:  # pylint: disable=broad-except
        pass
    return None


def _resolve_model(repo_id: str) -> tuple[str, str, bool]:
    """Return (model_type, model_args, is_gguf) for lm_eval CLI.

    Accepts:
      - llmpm repo ID:       "meta-llama/Llama-3.2-3B-Instruct"
      - lm_eval model_args:  "pretrained=gpt2"   (passed through as-is)

    lm_eval's ``gguf`` model type is an API client — it cannot load local
    .gguf files directly. For GGUF models we return is_gguf=True so the
    caller can start ``llmpm serve`` first and point lm_eval at the server.
    """
    # If the caller supplied a raw key=value string, pass it through unchanged.
    if "=" in repo_id and "/" not in repo_id.split("=")[0]:
        return "hf", repo_id, False

    # Look the model up in the llmpm registry.
    entry = registry.get_model(repo_id)
    if entry is None:
        # Not installed locally — treat as a HuggingFace Hub model ID.
        display.info(
            f"[bold]{repo_id}[/] not found in local registry — "
            "running directly from HuggingFace Hub."
        )
        return "hf", f"pretrained={repo_id}", False

    model_path = entry.get("path", "")
    files: list[str] = entry.get("files", [])

    gguf_files = [f for f in files if f.endswith(".gguf")]
    if gguf_files:
        return "local-chat-completions", "", True

    # Append max_length so lm_eval won't request more tokens than the model supports.
    max_len = _read_model_max_length(model_path)
    model_args = f"pretrained={model_path}"
    if max_len:
        model_args += f",max_length={max_len}"

    return "hf", model_args, False


def _wait_for_server(url: str, timeout: int = 60) -> bool:
    """Poll url/health until the server responds or timeout is reached."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            r = requests.get(f"{url}/health", timeout=2)
            if r.status_code == 200:
                return True
        except requests.exceptions.RequestException:
            pass
        time.sleep(1)
    return False


def _benchmark_deps() -> list[str]:
    """Read benchmark optional-deps from the installed llmpm package metadata."""
    from importlib.metadata import requires  # pylint: disable=import-outside-toplevel
    reqs = requires("llmpm") or []
    return [
        r.split(";")[0].strip()
        for r in reqs
        if 'extra == "benchmark"' in r
    ]


def _deps_hash() -> str:
    """Stable hash of benchmark deps — used to detect dependency changes."""
    import hashlib  # pylint: disable=import-outside-toplevel
    deps = "\n".join(sorted(_benchmark_deps()))
    return hashlib.sha256(deps.encode()).hexdigest()[:16]


def _auto_device() -> str:
    """Return the best available PyTorch device without importing torch."""
    import platform
    import subprocess as _sp
    if platform.system() == "Darwin":
        # Apple Silicon: prefer mps, fall back to cpu
        import shutil as _shutil
        _py = _shutil.which("python") or _shutil.which("python3") or "python"
        try:
            out = _sp.run(
                [_py, "-c",
                 "import torch; print('mps' if torch.backends.mps.is_available() else 'cpu')"],
                capture_output=True, text=True, timeout=10,
            )
            return out.stdout.strip() or "cpu"
        except Exception:  # pylint: disable=broad-except
            return "cpu"
    return "cpu"


def _ensure_benchmark_env() -> str:
    """Return path to lm_eval binary, auto-creating/rebuilding its venv as needed."""
    import venv as _venv
    from llmpm.core.registry import llmpm_home  # pylint: disable=import-outside-toplevel

    bench_venv = llmpm_home() / "benchmark_venv"
    lm_eval_bin = bench_venv / "bin" / "lm_eval"
    hash_file = bench_venv / ".deps_hash"
    current_hash = _deps_hash()

    needs_install = (
        not lm_eval_bin.exists()
        or not hash_file.exists()
        or hash_file.read_text().strip() != current_hash
    )

    if needs_install:
        if lm_eval_bin.exists():
            display.info("Benchmark dependencies changed — rebuilding environment…")
        else:
            display.info("First run: setting up benchmark environment…")

        _venv.create(str(bench_venv), with_pip=True, clear=True)

        pip = str(bench_venv / "bin" / "pip")
        deps = _benchmark_deps()
        display.info("Installing benchmark dependencies (this may take a few minutes)…")
        ret = subprocess.run(
            [pip, "install", *deps],
            check=False,
        ).returncode
        if ret != 0:
            display.error(
                "Failed to install benchmark dependencies.\n\n"
                f"  Try manually:  [bold]{pip} install {' '.join(deps)}[/]"
            )
            raise SystemExit(1)

        hash_file.write_text(current_hash)
        display.info("Benchmark environment ready.")
        display.blank()

    return str(lm_eval_bin), str(bench_venv)


def _stream_subprocess(cmd: list[str], extra_env: dict | None = None) -> tuple[int, list[str]]:
    """Run cmd, streaming stdout+stderr line-by-line. Returns (returncode, captured_lines)."""
    env = {**os.environ, "PYTHONUNBUFFERED": "1", **(extra_env or {})}
    captured: list[str] = []
    with subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
        text=True,
        env=env,
    ) as proc:
        assert proc.stdout is not None
        for line in proc.stdout:
            print(line, end="", flush=True)
            captured.append(line)
        proc.wait()
        return proc.returncode, captured


def _missing_extras(lines: list[str]) -> set[str]:
    """Scan captured output for lm-eval[X] install hints and return the extra names."""
    extras: set[str] = set()
    for line in lines:
        for match in _MISSING_EXTRA_RE.finditer(line):
            extras.add(match.group(1))
    return extras


def _install_extras(extras: set[str], bench_venv: str) -> None:
    """Install missing lm-eval extras into the benchmark venv."""
    pip = str(Path(bench_venv) / "bin" / "pip")
    packages = [f"lm-eval[{e}]" for e in sorted(extras)]
    display.info(f"Auto-installing missing extras: {', '.join(packages)}")
    ret = subprocess.run([pip, "install", *packages], check=False).returncode
    if ret != 0:
        display.warn(f"Failed to install {', '.join(packages)} — benchmark may still fail.")


def _find_results_json(output_dir: Path) -> Path | None:
    """Recursively find the first results*.json lm_eval wrote."""
    for p in sorted(output_dir.rglob("results*.json")):
        return p
    return None


def _generate_html_report(
    results_json: Path,
    repo_id: str,
    task_list: list[str],
    report_path: Path,
) -> None:
    """Parse lm_eval results JSON and write a self-contained HTML report."""
    with open(results_json, encoding="utf-8") as fh:
        data = json.load(fh)

    results: dict = data.get("results", {})
    config: dict = data.get("config", {})
    run_date: str = data.get("date", datetime.now().isoformat())

    # ── build task rows ───────────────────────────────────────────────────────
    task_rows_html = ""
    for task_name, metrics in results.items():
        desc = KNOWN_TASKS.get(task_name, "—")
        metric_cells = ""
        for key, val in metrics.items():
            if key.endswith(",none") and "stderr" not in key:
                label = key.replace(",none", "")
                pct = f"{val * 100:.2f}%" if isinstance(val, float) and val <= 1 else (
                    f"{val:.4f}" if isinstance(val, float) else str(val)
                )
                stderr_key = f"{label}_stderr,none"
                stderr_val = metrics.get(stderr_key, "")
                stderr_str = (
                    f" <span class='stderr'>± {stderr_val * 100:.2f}%</span>"
                    if isinstance(stderr_val, float)
                    else ""
                )
                metric_cells += f"<td class='metric'>{pct}{stderr_str}</td><th class='mlabel'>{label}</th>"
        task_rows_html += f"""
        <tr>
          <td class='task'><strong>{task_name}</strong><br><span class='desc'>{desc}</span></td>
          {metric_cells}
        </tr>"""

    # ── config table ──────────────────────────────────────────────────────────
    config_rows = ""
    cfg_display = {
        "Model": repo_id,
        "Backend": config.get("model", "—"),
        "Model args": config.get("model_args", "—"),
        "Tasks": ", ".join(task_list),
        "Batch size": config.get("batch_size", "—"),
        "Few-shot": config.get("num_fewshot", "task default"),
        "Limit": config.get("limit", "none"),
        "Device": config.get("device", "auto"),
        "Run date": run_date,
    }
    for k, v in cfg_display.items():
        config_rows += f"<tr><th>{k}</th><td>{v}</td></tr>"

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>llmpm benchmark — {repo_id}</title>
  <style>
    :root {{
      --bg: #0f1117; --surface: #1a1d27; --border: #2d3045;
      --text: #e2e8f0; --muted: #8892a4; --accent: #4ade80;
      --warn: #facc15; --blue: #60a5fa;
    }}
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{ background: var(--bg); color: var(--text); font-family: system-ui, sans-serif; padding: 2rem; }}
    h1 {{ font-size: 1.6rem; font-weight: 700; margin-bottom: .25rem; }}
    h1 span {{ color: var(--accent); }}
    .subtitle {{ color: var(--muted); font-size: .9rem; margin-bottom: 2rem; }}
    section {{ margin-bottom: 2.5rem; }}
    h2 {{ font-size: 1rem; text-transform: uppercase; letter-spacing: .08em;
          color: var(--muted); margin-bottom: .75rem; border-bottom: 1px solid var(--border); padding-bottom: .4rem; }}
    table {{ width: 100%; border-collapse: collapse; font-size: .92rem; }}
    th, td {{ padding: .55rem .8rem; text-align: left; border-bottom: 1px solid var(--border); }}
    thead th {{ background: var(--surface); color: var(--blue); font-weight: 600; }}
    tr:hover td, tr:hover th {{ background: #ffffff08; }}
    td.task {{ min-width: 200px; }}
    .desc {{ color: var(--muted); font-size: .8rem; margin-top: .2rem; }}
    td.metric {{ font-size: 1rem; font-weight: 700; color: var(--accent); text-align: right; }}
    th.mlabel {{ color: var(--muted); font-size: .78rem; font-weight: 400; }}
    .stderr {{ color: var(--muted); font-size: .8rem; font-weight: 400; }}
    .config-table th {{ color: var(--muted); font-weight: 500; width: 140px; }}
    .config-table td {{ font-family: monospace; font-size: .85rem; word-break: break-all; }}
    .badge {{ display: inline-block; background: var(--surface); border: 1px solid var(--border);
              border-radius: 4px; padding: .1rem .5rem; font-size: .78rem; color: var(--muted);
              margin-right: .4rem; }}
    footer {{ color: var(--muted); font-size: .78rem; margin-top: 3rem; text-align: center; }}
  </style>
</head>
<body>
  <h1>llmpm <span>benchmark</span></h1>
  <p class="subtitle">
    <span class="badge">{repo_id}</span>
    <span class="badge">{", ".join(task_list)}</span>
    <span class="badge">{run_date}</span>
  </p>

  <section>
    <h2>Results</h2>
    <table>
      <thead>
        <tr><th>Task</th><th colspan="10">Metrics</th></tr>
      </thead>
      <tbody>
        {task_rows_html}
      </tbody>
    </table>
  </section>

  <section>
    <h2>Run Configuration</h2>
    <table class="config-table">
      <tbody>{config_rows}</tbody>
    </table>
  </section>

  <footer>Generated by <strong>llmpm benchmark</strong> · powered by <a href="https://github.com/EleutherAI/lm-evaluation-harness" style="color:var(--blue)">lm-evaluation-harness</a></footer>
</body>
</html>
"""
    report_path.write_text(html, encoding="utf-8")


def _print_task_list() -> None:
    from rich.table import Table
    from rich.console import Console
    from rich import box

    console = Console()
    table = Table(
        title="Supported Benchmarks",
        box=box.SIMPLE_HEAD,
        show_header=True,
        header_style="bold cyan",
        min_width=80,
    )
    table.add_column("Task / Group", style="bold", min_width=18)
    table.add_column("Description")
    for name, desc in KNOWN_TASKS.items():
        table.add_row(name, desc)
    console.print()
    console.print(table)
    console.print(
        "  [dim]Any task name from lm_eval/tasks/ is also valid — "
        "the list above covers the most common ones.[/dim]\n"
    )


@click.command("benchmark")
@click.argument("repo_id", metavar="MODEL", required=False, default=None)
@click.option(
    "--tasks", "-t",
    required=False,
    default=None,
    help="Comma-separated benchmark tasks/groups (e.g. ifeval,hellaswag,mmlu).",
)
@click.option(
    "--num-fewshot", "-n",
    default=None,
    type=int,
    help="Override few-shot count for all tasks (uses task default when omitted).",
)
@click.option(
    "--limit", "-l",
    default=None,
    type=float,
    help="Limit examples per task. Integer = count, <1.0 = fraction of dataset.",
)
@click.option(
    "--batch-size", "-b",
    default="auto",
    show_default=True,
    help="Inference batch size ('auto' lets lm_eval choose).",
)
@click.option(
    "--device",
    default=None,
    help="PyTorch device (e.g. 'cpu', 'cuda:0'). Auto-detected when omitted.",
)
@click.option(
    "--output", "-o",
    default=None,
    type=click.Path(dir_okay=True, writable=True),
    help="Directory to write lm_eval result JSON files into.",
)
@click.option(
    "--list-tasks",
    is_flag=True,
    default=False,
    help="Print supported benchmark tasks and exit.",
)
def benchmark(
    repo_id: str | None,
    tasks: str | None,
    num_fewshot: int | None,
    limit: float | None,
    batch_size: str,
    device: str | None,
    output: str | None,
    list_tasks: bool,
) -> None:
    """Run LLM evaluation benchmarks against an installed model.

    MODEL can be an llmpm repo ID (e.g. meta-llama/Llama-3.2-3B-Instruct)
    or a raw lm_eval model_args string (e.g. pretrained=gpt2).

    \b
    Examples:
      llmpm benchmark meta-llama/Llama-3.2-3B-Instruct --tasks ifeval,hellaswag
      llmpm benchmark pretrained=gpt2 --tasks mmlu --num-fewshot 5
      llmpm benchmark meta-llama/Llama-3.2-3B-Instruct --tasks openllm
      llmpm benchmark unsloth/Llama-GGUF --tasks hellaswag --limit 100
      llmpm benchmark --list-tasks
    """
    if list_tasks:
        _print_task_list()
        return

    # ── validate required args ────────────────────────────────────────────────
    if not repo_id:
        display.error(
            "MODEL is required.\n\n"
            "  Usage:  [bold cyan]llmpm benchmark <model> --tasks <task,...>[/]\n"
            "  Hint:   [bold cyan]llmpm benchmark --list-tasks[/]  to see supported tasks."
        )
        raise SystemExit(1)

    if not tasks:
        display.error(
            "--tasks is required.\n\n"
            "  Example:  [bold cyan]llmpm benchmark {repo_id} --tasks ifeval,hellaswag[/]\n"
            "  Hint:     [bold cyan]llmpm benchmark --list-tasks[/]"
        )
        raise SystemExit(1)

    # ── ensure benchmark venv + lm_eval binary ────────────────────────────────
    lm_eval_bin, bench_venv = _ensure_benchmark_env()
    lm_eval_cmd = [lm_eval_bin]
    # Use a venv-local HF datasets cache to avoid stale metadata from the
    # system cache (written by a different datasets version).
    bench_env = {"HF_DATASETS_CACHE": str(Path(bench_venv) / "hf_cache")}

    # ── resolve model ─────────────────────────────────────────────────────────
    # Try fuzzy-matching against installed models first.
    if "=" not in repo_id:
        matches = registry.find_models(repo_id)
        if matches and len(matches) > 1:
            result = pick_installed_model(matches, repo_id)
            if result is None:
                raise SystemExit(0)
            repo_id = result[0]
        elif matches and len(matches) == 1:
            matched_id = matches[0][0]
            if matched_id != repo_id:
                display.info(f"Using [bold]{matched_id}[/]")
            repo_id = matched_id

    model_type, model_args, is_gguf = _resolve_model(repo_id)
    task_list = [t.strip() for t in tasks.split(",") if t.strip()]

    # Replace any tasks that pull gated datasets with their open equivalents.
    for i, task in enumerate(task_list):
        if task in _GATED_SUBSTITUTIONS:
            replacement = _GATED_SUBSTITUTIONS[task]
            display.warn(
                f"[bold]{task}[/] includes gated datasets — "
                f"auto-substituting with [bold]{replacement}[/]"
            )
            task_list[i] = replacement

    # ── print header ──────────────────────────────────────────────────────────
    display.header(f"benchmark  [dim]{repo_id}[/]")
    display.blank()
    display.info(f"Backend    : [bold]{'gguf (via llmpm serve)' if is_gguf else model_type}[/]")
    display.info(f"Tasks      : [bold]{', '.join(task_list)}[/]")
    if num_fewshot is not None:
        display.info(f"Few-shot   : [bold]{num_fewshot}[/]")
    if limit is not None:
        display.info(f"Limit      : [bold]{limit}[/]")
    display.blank()

    # ── GGUF: auto-start llmpm serve, point lm_eval at the local API ──────────
    serve_proc: subprocess.Popen | None = None
    serve_url = "http://localhost:8080"

    if is_gguf:
        import shutil as _shutil
        llmpm_bin = _shutil.which("llmpm") or sys.argv[0]
        display.info("GGUF model detected — starting [bold]llmpm serve[/] in background…")
        serve_proc = subprocess.Popen(
            [llmpm_bin, "serve", repo_id, "--port", "8080"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        display.info("Waiting for server to be ready…")
        if not _wait_for_server(serve_url, timeout=120):
            serve_proc.terminate()
            display.error(
                "Server did not start within 120 s.\n\n"
                "  Start it manually first:  [bold cyan]llmpm serve {repo_id}[/]\n"
                "  Then re-run with:          [bold cyan]--model-url http://localhost:8080[/]"
            )
            raise SystemExit(1)
        display.info(f"Server ready at [bold]{serve_url}[/]")
        display.blank()
        model_type = "local-chat-completions"
        model_args = f"model={repo_id},base_url={serve_url}/v1/"

    # ── set up output directory ───────────────────────────────────────────────
    _tmpdir = None
    if output:
        out_dir = Path(output)
        out_dir.mkdir(parents=True, exist_ok=True)
    else:
        _tmpdir = tempfile.mkdtemp(prefix="llmpm_bench_")
        out_dir = Path(_tmpdir)

    # ── build lm_eval run command ─────────────────────────────────────────────
    cmd: list[str] = [
        *lm_eval_cmd,
        "--model", model_type,
        "--model_args", model_args,
        "--tasks", ",".join(task_list),
        "--batch_size", str(batch_size),
        "--output_path", str(out_dir),
    ]
    if num_fewshot is not None:
        cmd += ["--num_fewshot", str(num_fewshot)]
    if limit is not None:
        cmd += ["--limit", str(limit)]
    if not is_gguf:
        cmd += ["--device", device or _auto_device()]
    if any(t in _BUNDLED_TASK_NAMES for t in task_list):
        cmd += ["--include_path", str(_BUNDLED_TASKS_DIR)]

    # ── run lm_eval, streaming output line-by-line ────────────────────────────
    returncode = 0
    try:
        returncode, captured = _stream_subprocess(cmd, extra_env=bench_env)

        # Auto-install missing lm-eval extras and retry once.
        if returncode != 0:
            extras = _missing_extras(captured)
            if extras:
                display.blank()
                _install_extras(extras, bench_venv)
                display.info("Retrying benchmark…")
                display.blank()
                returncode, captured = _stream_subprocess(cmd, extra_env=bench_env)

    except KeyboardInterrupt:
        display.warn("Benchmark interrupted.")
        returncode = 130
    finally:
        if serve_proc is not None:
            serve_proc.terminate()
            display.info("llmpm serve stopped.")

    # ── generate HTML report ──────────────────────────────────────────────────
    if returncode == 0:
        results_json = _find_results_json(out_dir)
        if results_json:
            _model_slug = repo_id.replace("/", "_").replace("=", "_")[:40]
            _tasks_slug = "_".join(task_list)[:40]
            _ts = datetime.now().strftime("%Y%m%d_%H%M%S")
            report_name = f"report_{_model_slug}_{_tasks_slug}_{_ts}.html"
            report_path = Path(output or ".") / report_name
            _generate_html_report(results_json, repo_id, task_list, report_path)
            display.blank()
            display.info(f"Report saved → [bold]{report_path.resolve()}[/]")
        else:
            display.warn("Could not find lm_eval results JSON — skipping HTML report.")
    elif returncode != 130:
        display.error(
            f"lm_eval exited with code {returncode}.\n\n"
            "  Check task names with:  [bold cyan]llmpm benchmark --list-tasks[/]"
        )

    raise SystemExit(returncode if returncode != 0 else None)
