# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "openai>=1.60",
#   "huggingface_hub>=0.28",
#   "matplotlib>=3.8",
# ]
# ///
"""
HF Jobs entrypoint for the InferGauge benchmark.

Runs entirely on Hugging Face infrastructure. Pulls the harness out of the Space
repo, runs the benchmark against Inference Providers, then pushes leaderboard.csv,
results.csv, report.md and the charts back into the same Space — which makes the
public leaderboard update itself.

Run it once:

    hf jobs uv run --with-secrets HF_TOKEN \
      https://huggingface.co/spaces/Infergauge-load/llm-performance-and-load/raw/main/job.py

Run it monthly:

    hf jobs scheduled uv run "13 7 1 * *" --with-secrets HF_TOKEN \
      https://huggingface.co/spaces/Infergauge-load/llm-performance-and-load/raw/main/job.py

Environment:
    HF_TOKEN   required, WRITE scope (it both calls inference and pushes results)
    IG_SMOKE   "1" for the 3-model x 5-prompt smoke run (default). "0" for the full run.
    IG_REPEATS repeats per prompt on a full run (default 3)
    IG_JUDGE   optional judge model id, e.g. Qwen/Qwen3-8B. Blank = no judge.
    IG_CONCURRENCY  parallel requests per model, 1-25. Default 1 (single-stream).
    IG_MODELS  optional comma-separated model ids to restrict the run to. Use with
               IG_NO_PUSH=1 to trial a candidate model without touching the Space.
    IG_LIMIT_PROMPTS  optional cap on the number of prompts, for a quick trial.
    IG_NO_PUSH "1" runs the benchmark and prints the leaderboard to the log without
               publishing anything. The Space is public and the weekly history is
               cumulative, so a trial must never overwrite either.
    IG_REPO_ID override the Space results are pushed to. Defaults to the one above.
"""

from __future__ import annotations

import csv
import io
import json
import os
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

REPO_ID = os.environ.get("IG_REPO_ID", "Infergauge-load/llm-performance-and-load")
REPO_TYPE = "space"

HARNESS = [
    "benchmark.py",
    "quality.py",
    "hf_client.py",
    "prompts.json",
    "pricing.json",
    "slo.json",
    "model_shortlist.csv",
]

WORK = Path("/tmp/ig")


def fetch_harness(token: str) -> Path:
    """Pull the harness out of the Space repo into a local working layout."""
    (WORK / "integration").mkdir(parents=True, exist_ok=True)
    (WORK / "testing").mkdir(parents=True, exist_ok=True)

    # benchmark.py expects: <root>/testing/*.py and <root>/integration/{hf_client,model_shortlist}
    placement = {
        "benchmark.py": WORK / "testing" / "benchmark.py",
        "quality.py": WORK / "testing" / "quality.py",
        "prompts.json": WORK / "testing" / "prompts.json",
        "pricing.json": WORK / "testing" / "pricing.json",
        "slo.json": WORK / "testing" / "slo.json",
        "hf_client.py": WORK / "integration" / "hf_client.py",
        "model_shortlist.csv": WORK / "integration" / "model_shortlist.csv",
    }

    # slo.json is optional: benchmark.py falls back to built-in thresholds, so a
    # Space that hasn't been updated yet still runs instead of dying on a 404.
    optional = {"slo.json"}

    for remote in HARNESS:
        try:
            local = hf_hub_download(
                repo_id=REPO_ID, repo_type=REPO_TYPE, filename=remote, token=token
            )
        except Exception:  # noqa: BLE001
            if remote in optional:
                print(f"  -- {remote} not in repo, using defaults", flush=True)
                continue
            raise
        placement[remote].write_bytes(Path(local).read_bytes())
        print(f"  fetched {remote}", flush=True)
    return WORK


def run_benchmark(root: Path) -> int:
    smoke = os.environ.get("IG_SMOKE", "1") != "0"
    cmd = [sys.executable, str(root / "testing" / "benchmark.py"), "--sleep", "0.4"]
    cmd += ["--smoke"] if smoke else ["--repeats", os.environ.get("IG_REPEATS", "3")]
    conc = os.environ.get("IG_CONCURRENCY", "1").strip()
    if conc and conc != "1":
        cmd += ["--concurrency", conc]
    judge = os.environ.get("IG_JUDGE", "").strip()
    if judge:
        cmd += ["--judge-model", judge]

    # A trial set, so a new model can be measured before it joins the weekly run.
    # benchmark.py already takes --models; this only plumbs it through to the job.
    models = os.environ.get("IG_MODELS", "").strip()
    if models:
        cmd += ["--models", models]
    prompts = os.environ.get("IG_LIMIT_PROMPTS", "").strip()
    if prompts:
        cmd += ["--limit-prompts", prompts]

    print(f"\n$ {' '.join(cmd)}\n", flush=True)
    return subprocess.call(cmd, cwd=str(root))


def error_rate(root: Path) -> tuple[float, int, int]:
    """Share of requests that failed, from results.csv."""
    import csv

    path = root / "testing" / "results.csv"
    if not path.exists():
        return 1.0, 0, 0
    with path.open() as f:
        rows = list(csv.DictReader(f))
    if not rows:
        return 1.0, 0, 0
    bad = sum(1 for r in rows if r.get("error"))
    return bad / len(rows), bad, len(rows)


def push_results(root: Path, token: str) -> None:
    """Upload results back into the Space. leaderboard.csv goes to the root, where
    index.html reads it; everything else lands under results/.

    Refuses to publish a mostly-failed run. The Space is public: a leaderboard of
    error rows is worse than the honest 'not run yet' state it would replace. Set
    IG_FORCE_PUSH=1 to override when you genuinely want the failures published.
    """
    api = HfApi(token=token)
    testing = root / "testing"

    rate, bad, total = error_rate(root)
    threshold = float(os.environ.get("IG_MAX_ERROR_RATE", "0.5"))
    if rate > threshold and os.environ.get("IG_FORCE_PUSH") != "1":
        print(f"!! {bad}/{total} requests failed ({rate:.0%}) — refusing to publish.",
              flush=True)
        print("   The Space is public; a leaderboard of errors is worse than the "
              "'not run yet' state.", flush=True)
        print("   Fix the cause, then re-run. Set IG_FORCE_PUSH=1 to publish anyway.",
              flush=True)
        print(f"   Local results kept at {testing}", flush=True)
        return

    lb = testing / "leaderboard.csv"
    if not lb.exists():
        print("!! no leaderboard.csv produced — nothing to push", flush=True)
        return

    api.upload_file(
        path_or_fileobj=str(lb),
        path_in_repo="leaderboard.csv",
        repo_id=REPO_ID,
        repo_type=REPO_TYPE,
        commit_message="Benchmark run: refresh leaderboard",
    )
    print("  pushed leaderboard.csv -> repo root", flush=True)

    # Keep a per-concurrency copy as well, so a new run never destroys the previous
    # level's data. leaderboard.csv is always the most recent run.
    conc = os.environ.get("IG_CONCURRENCY", "1").strip() or "1"
    api.upload_file(
        path_or_fileobj=str(lb),
        path_in_repo=f"leaderboard-c{conc}.csv",
        repo_id=REPO_ID,
        repo_type=REPO_TYPE,
        commit_message=f"Benchmark run: leaderboard at concurrency {conc}",
    )
    print(f"  pushed leaderboard-c{conc}.csv", flush=True)

    extras = {
        "results.csv": testing / "results.csv",
        "report.md": testing / "report.md",
    }
    for name, path in extras.items():
        if path.exists():
            api.upload_file(
                path_or_fileobj=str(path),
                path_in_repo=f"results/{name}",
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                commit_message=f"Benchmark run: {name}",
            )
            print(f"  pushed {name}", flush=True)

    charts = testing / "charts"
    if charts.is_dir() and any(charts.iterdir()):
        api.upload_folder(
            folder_path=str(charts),
            path_in_repo="results/charts",
            repo_id=REPO_ID,
            repo_type=REPO_TYPE,
            commit_message="Benchmark run: charts",
        )
        print("  pushed charts", flush=True)

    # Archiving, the manifest and the weekly post all come last, and none of them can
    # take the run down with them. The results are already in the Space by this point;
    # losing the history entry for one week is recoverable, losing the results is not.
    try:
        date = run_date()
        archived = archive_run(api, lb, conc, date)
        runs = update_manifest(api, token, conc, date, archived)
    except Exception as e:  # noqa: BLE001
        print(f"  !! archive/manifest failed ({type(e).__name__}: {e}) — "
              f"results are pushed, history entry missing", flush=True)
        return

    try:
        announce(api, token, date, runs)
    except Exception as e:  # noqa: BLE001
        print(f"  !! announcement failed ({type(e).__name__}: {e}) — results are pushed",
              flush=True)


def run_date() -> str:
    """UTC date stamp shared by every level of one weekly set."""
    return os.environ.get("IG_RUN_DATE") or datetime.now(timezone.utc).strftime("%Y-%m-%d")


def archive_run(api, lb: Path, conc: str, date: str) -> str:
    """Keep every run forever under its own date.

    leaderboard-c<N>.csv is 'latest' and gets overwritten. These do not: the point
    of running weekly is the history, and a file that is rewritten every Monday has
    no history in it.
    """
    path = f"runs/{date}-c{conc}.csv"
    api.upload_file(
        path_or_fileobj=str(lb), path_in_repo=path,
        repo_id=REPO_ID, repo_type=REPO_TYPE,
        commit_message=f"Archive run {date} at concurrency {conc}",
    )
    print(f"  archived {path}", flush=True)
    return path


def update_manifest(api, token: str, conc: str, date: str, path: str) -> list:
    """runs/manifest.json lists every archived run. index.html reads it to build
    the dated sections, so it is the index of the whole history."""
    runs = []
    try:
        local = hf_hub_download(repo_id=REPO_ID, repo_type=REPO_TYPE,
                                filename="runs/manifest.json", token=token)
        runs = json.loads(Path(local).read_text()).get("runs", [])
    except Exception:  # noqa: BLE001 — first run, no manifest yet
        pass

    runs = [r for r in runs if not (r.get("date") == date and str(r.get("concurrency")) == str(conc))]
    runs.append({"date": date, "concurrency": int(conc), "file": path})
    runs.sort(key=lambda r: (r["date"], r["concurrency"]), reverse=True)

    api.upload_file(
        path_or_fileobj=json.dumps({"runs": runs}, indent=1).encode(),
        path_in_repo="runs/manifest.json",
        repo_id=REPO_ID, repo_type=REPO_TYPE,
        commit_message=f"Manifest: add {date} c{conc}",
    )
    print(f"  manifest now lists {len(runs)} run(s)", flush=True)
    return runs


def _read_remote_csv(token: str, path: str) -> list:
    try:
        local = hf_hub_download(repo_id=REPO_ID, repo_type=REPO_TYPE,
                                filename=path, token=token)
        return list(csv.DictReader(io.StringIO(Path(local).read_text())))
    except Exception:  # noqa: BLE001
        return []


def _f(row, key):
    try:
        return float(row.get(key) or "")
    except ValueError:
        return None


# The four views below mirror the Overview / Latency / Quality / Tokens & cost
# tabs in index.html. Markdown has neither tabs nor a column-group header row,
# so each view becomes a <details> block and every per-level column carries an
# @<level> suffix instead of sitting under a spanning header.
_VIEWS = (
    # title, open by default, flat columns, per-level columns
    ("Overview", True,
     (("Class", "size_class", None), ("Provider", "provider", None),
      ("$/1M out", "cost_per_1m_output_usd", 2)),
     (("TTFT p95", "ttft_p95_s", 2), ("Err %", "error_rate_pct", 1))),
    ("Latency", False,
     (("Provider", "provider", None),),
     (("TTFT p95", "ttft_p95_s", 2), ("ITL p95 ms", "itl_p95_ms", 0),
      ("Tok/s", "tokens_per_s", 0))),
    ("Quality", False,
     (("Class", "size_class", None), ("Licence", "license", None)),
     (("Quality", "quality_score", 2), ("Goodput %", "goodput_pct", 0),
      ("Err %", "error_rate_pct", 1))),
    ("Tokens & cost", False,
     (("Provider", "provider", None), ("$/1M out", "cost_per_1m_output_usd", 2)),
     (("In tok", "in_tokens_avg", 0), ("Out tok", "out_tokens_avg", 0),
      ("Tok/s", "tokens_per_s", 0))),
)

DASH = "—"


def _cell(row, key, d):
    """One table cell. A missing model, a missing column and an unparseable
    number all render as an em dash rather than a 0 or a blank."""
    if not row:
        return DASH
    raw = row.get(key)
    if raw is None or raw == "":
        return DASH
    if d is None:
        return str(raw)
    v = _f(row, key)
    return DASH if v is None else f"{v:.{d}f}"


def view_tables(token: str, todays: list) -> list:
    """The leaderboard's four views, as markdown, one <details> per view.

    Driven entirely by the archived CSVs for this date, so the tables cannot
    drift from what the Space is serving. Rows are ordered the way the
    leaderboard opens: TTFT p95 at the highest level, ascending.
    """
    levels = [str(r["concurrency"]) for r in todays]
    by_level = {}
    for r in todays:
        rows = _read_remote_csv(token, r["file"])
        by_level[str(r["concurrency"])] = {x["model_id"]: x for x in rows if x.get("model_id")}

    models = sorted({m for lv in by_level.values() for m in lv})
    if not models:
        return []

    first, last = levels[0], levels[-1]
    models.sort(key=lambda m: (_f(by_level[last].get(m) or {}, "ttft_p95_s")
                               if _f(by_level[last].get(m) or {}, "ttft_p95_s") is not None
                               else float("inf"), m))

    def delta(m):
        a = _f(by_level[first].get(m) or {}, "ttft_p95_s")
        b = _f(by_level[last].get(m) or {}, "ttft_p95_s")
        if not a or b is None:
            return DASH
        return f"{(b / a - 1) * 100:+.0f}%"

    out = ["", f"Per-model figures, all {len(levels)} levels. `@1` means at one "
               f"concurrent user, `@{last}` at {last}."]
    for title, is_open, flat, per in _VIEWS:
        head = [c[0] for c in flat]
        for lv in levels:
            head += [f"{c[0]} @{lv}" for c in per]
        if title == "Overview":
            head += [f"Goodput % @{last}", f"Δ p95 1→{last}"]

        body = []
        for m in models:
            cells = [f"`{m}`"]
            base = next((by_level[lv][m] for lv in levels if m in by_level[lv]), None)
            cells += [_cell(base, k, d) for _, k, d in flat]
            for lv in levels:
                row = by_level[lv].get(m)
                cells += [_cell(row, k, d) for _, k, d in per]
            if title == "Overview":
                cells += [_cell(by_level[last].get(m), "goodput_pct", 0), delta(m)]
            body.append("| " + " | ".join(cells) + " |")

        out += [
            "",
            f"<details{' open' if is_open else ''}>",
            f"<summary><b>{title}</b></summary>",
            "",
            "| Model | " + " | ".join(head) + " |",
            "|---" * (len(head) + 1) + "|",
            *body,
            "",
            "</details>",
        ]
    return out


def announce(api, token: str, date: str, runs: list) -> None:
    """Post a new Community discussion for this week's set.

    A new post per week rather than edits to an old one, so each week's numbers stay
    readable on their own. Only fires when IG_ANNOUNCE=1, which is set on the LAST
    level of the weekly sequence — otherwise you would get one post per level.
    """
    if os.environ.get("IG_ANNOUNCE") != "1":
        return

    todays = sorted([r for r in runs if r.get("date") == date],
                    key=lambda r: r["concurrency"])
    if not todays:
        print("  nothing archived for today — no announcement", flush=True)
        return

    lines = [
        f"Weekly run, {date}. Same harness, same 20 prompts, one provider pinned per "
        f"model, no client-side retries. Levels measured sequentially, never in parallel.",
        "",
        "| Concurrency | Models >=90% goodput | Models with errors | TTFT p95 range |",
        "|---|---|---|---|",
    ]
    detail = []
    for r in todays:
        rows = _read_remote_csv(token, r["file"])
        if not rows:
            continue
        gp = [_f(x, "goodput_pct") for x in rows]
        p95 = [v for v in (_f(x, "ttft_p95_s") for x in rows) if v]
        errs = [x for x in rows if (_f(x, "error_rate_pct") or 0) > 0]
        good = sum(1 for v in gp if v is not None and v >= 90)
        spread = f"{min(p95):.2f}-{max(p95):.2f}s ({max(p95)/min(p95):.0f}x)" if p95 else "-"
        lines.append(f"| {r['concurrency']} | {good} / {len(rows)} | {len(errs)} | {spread} |")
        if errs:
            worst = max(errs, key=lambda x: _f(x, "error_rate_pct") or 0)
            detail.append(f"- At {r['concurrency']}: worst error rate was "
                          f"{worst['model_id']} at {worst['error_rate_pct']}%.")

    if detail:
        lines += ["", "Notable:", *detail]

    lines += view_tables(token, todays)

    lines += [
        "",
        "Goodput counts a request only if it met every threshold at once: first token "
        "under 2s, finished under 15s, inter-token latency p95 under 80ms, format score "
        "at or above 0.8, no error. Thresholds are in `slo.json`.",
        "",
        "These are shared endpoints carrying load nobody here controls, and each figure "
        "is 60 requests. Week-to-week movement of a few points is noise, not signal. "
        "The per-run CSVs are archived under `runs/` so you can check the trend yourself.",
        "",
        f"Full matrix and method: https://huggingface.co/spaces/{REPO_ID}",
    ]

    d = api.create_discussion(
        repo_id=REPO_ID, repo_type=REPO_TYPE,
        title=f"Weekly run {date}: {', '.join(str(r['concurrency']) for r in todays)} concurrent users",
        description="\n".join(lines),
        token=token,
    )
    print(f"  posted discussion #{getattr(d, 'num', '?')}", flush=True)


def error_digest(root: Path, max_msgs: int = 3) -> None:
    """Per-model error counts and the distinct messages behind them, into the log.

    results.csv carries the exception strings, but it is only uploaded on a real
    run and a trial deliberately uploads nothing — so a 100% error rate used to
    arrive with no way to find out why, and the container is gone by the time
    anyone asks. Printing a digest costs nothing and is the difference between
    "ovhcloud failed" and "ovhcloud rejected this field".

    Wrapped end to end: a reporting bug must never take down a run whose
    results are otherwise fine.
    """
    import csv
    from collections import defaultdict
    try:
        path = root / "testing" / "results.csv"
        if not path.exists():
            return
        with path.open() as f:
            rows = list(csv.DictReader(f))
        total = defaultdict(int)
        errs = defaultdict(list)
        for r in rows:
            m = r.get("model_id") or "?"
            total[m] += 1
            e = (r.get("error") or "").strip()
            if e:
                errs[m].append(e)
        if not errs:
            return
        print("\n----- errors by model -----", flush=True)
        for m in sorted(errs, key=lambda k: -len(errs[k])):
            prov = next((r.get("provider", "") for r in rows if r.get("model_id") == m), "")
            print(f"  {m} ({prov}): {len(errs[m])}/{total[m]} failed", flush=True)
            for msg in list(dict.fromkeys(errs[m]))[:max_msgs]:
                print(f"      {msg[:300]}", flush=True)
        print("----- end -----\n", flush=True)
    except Exception as e:  # noqa: BLE001
        print(f"  !! error digest failed ({type(e).__name__}: {e}) — results unaffected",
              flush=True)


def show_local(root: Path) -> None:
    """Print the leaderboard into the job log.

    A trial run publishes nothing, so the log is the only place its numbers
    appear — the container is gone by the time anyone looks at the Space.
    """
    lb = root / "testing" / "leaderboard.csv"
    if not lb.exists():
        print("  no leaderboard.csv produced", flush=True)
        return
    print("\n----- leaderboard.csv -----", flush=True)
    print(lb.read_text().rstrip(), flush=True)
    print("----- end -----\n", flush=True)


def main() -> int:
    token = os.environ.get("HF_TOKEN")
    if not token:
        print("HF_TOKEN is not set. Pass it with --with-secrets HF_TOKEN.", file=sys.stderr)
        return 2

    mode = "SMOKE (3 models x 5 prompts)" if os.environ.get("IG_SMOKE", "1") != "0" \
        else f"FULL (repeats={os.environ.get('IG_REPEATS', '3')})"
    print(f"InferGauge benchmark job — {mode}\nrepo: {REPO_ID}\n", flush=True)

    print("Fetching harness…", flush=True)
    root = fetch_harness(token)

    rc = run_benchmark(root)
    print(f"\nbenchmark exited {rc}", flush=True)
    error_digest(root)

    if os.environ.get("IG_NO_PUSH") == "1":
        print("\nIG_NO_PUSH=1 — trial run. Nothing is published: the leaderboard,"
              "\n  the per-level copy, results/, runs/ and the manifest are all"
              "\n  left exactly as they are.", flush=True)
        show_local(root)
        return rc

    print("\nPushing results…", flush=True)
    push_results(root, token)

    print("\nDone. Leaderboard: "
          f"https://huggingface.co/spaces/{REPO_ID}", flush=True)
    return rc


if __name__ == "__main__":
    raise SystemExit(main())