#!/usr/bin/env python3
# Compile-time regression check.
#
# Times `ail check FILE` and `ail build --opt=-O0 FILE` over a curated
# corpus, drops the slowest run, takes the median of the rest, diffs
# against bench/baseline_compile.json. Exits 0 if every metric is within
# tolerance, 1 on any regression, 2 on infrastructure failure (missing
# fixture, ail spawn error).
#
# Methodology caveats this bench cannot work around:
#   - Wall time of `ail check` is dominated by subprocess spawn (~5-10ms
#     on Linux); typechecker work below that scale is invisible.
#   - Wall time of `ail build` is dominated by clang's link step
#     (~100ms+); the AILang IR-emit + codegen contribution is a
#     small slice of the total.
# Both are still useful as catastrophe detectors (10x slowdowns visible)
# even if subtler regressions need a profiler. The right tool for finer-
# grained codegen perf is the runtime bench (bench/run.sh), not this one.
#
# Usage:
#     bench/compile_check.py                  # run bench, exit 0/1/2
#     bench/compile_check.py -n 10            # more runs, tighter median
#     bench/compile_check.py --update-baseline
#     bench/compile_check.py --baseline path  # alternate baseline file

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_BASELINE = ROOT / "bench" / "baseline_compile.json"
AIL = ROOT / "target" / "release" / "ail"
EXAMPLES = ROOT / "examples"

# Curated corpus. Pinned by name so the baseline keys are stable across
# example-directory churn. Each entry must have a corresponding
# `<name>.ail.json` under examples/.
CORPUS = [
    # Surface coverage: each fixture exercises a different feature set.
    "hello",                    # IO baseline
    "list_map_poly",            # polymorphism + HOF
    "local_rec_capture",        # let-rec capture, closure path
    "borrow_own_demo",          # explicit modes
    "nested_pat",               # nested pattern matching
    # Bench fixtures (correlation with bench/check.py runtime baseline).
    "bench_list_sum",
    "bench_tree_walk",
    "bench_closure_chain",
    "bench_hof_pipeline",
    "bench_compute_intsum",     # 21'd; LLVM-folded at runtime, but compile path is normal
    "bench_compute_collatz",    # 21'd
    "bench_list_sum_explicit",  # 21'f; explicit-mode pair fixture
]


def time_one(args: list[str]) -> float:
    """Run a command, return wall-time in milliseconds. Raises if exit != 0."""
    t0 = time.monotonic()
    proc = subprocess.run(args, capture_output=True)
    t1 = time.monotonic()
    if proc.returncode != 0:
        raise RuntimeError(
            f"command failed (exit {proc.returncode}): {' '.join(args)}\n"
            f"stderr: {proc.stderr.decode(errors='replace')}"
        )
    return (t1 - t0) * 1000.0


def median_drop_slowest(runs: list[float]) -> float:
    """Drop the slowest run, return median of the rest."""
    if len(runs) < 2:
        return runs[0] if runs else 0.0
    kept = sorted(runs)[:-1]
    n = len(kept)
    if n % 2 == 1:
        return kept[n // 2]
    return 0.5 * (kept[n // 2 - 1] + kept[n // 2])


def measure(num_runs: int) -> dict[str, dict[str, float]]:
    """Return {section: {fixture: ms}}. Aborts on missing fixture / spawn fail."""
    if not AIL.is_file():
        print(f"missing release ail binary at {AIL}; building...", file=sys.stderr)
        subprocess.run(
            ["cargo", "build", "--release", "-p", "ail"],
            cwd=str(ROOT), check=True, capture_output=True,
        )

    out: dict[str, dict[str, float]] = {"check_ms": {}, "build_O0_ms": {}}
    with tempfile.NamedTemporaryFile(prefix="bench_compile_", delete=False) as tf:
        out_path = tf.name

    try:
        for fixture in CORPUS:
            src = EXAMPLES / f"{fixture}.ail.json"
            if not src.is_file():
                print(f"missing fixture: {src}", file=sys.stderr)
                sys.exit(2)

            try:
                check_runs = [time_one([str(AIL), "check", str(src)]) for _ in range(num_runs)]
                build_runs = [time_one([str(AIL), "build", "--opt=-O0",
                                        str(src), "-o", out_path]) for _ in range(num_runs)]
            except RuntimeError as e:
                print(f"spawn error on {fixture}: {e}", file=sys.stderr)
                sys.exit(2)

            out["check_ms"][fixture] = median_drop_slowest(check_runs)
            out["build_O0_ms"][fixture] = median_drop_slowest(build_runs)
            print(f"  {fixture:<32} check={out['check_ms'][fixture]:6.1f}ms  "
                  f"build={out['build_O0_ms'][fixture]:6.1f}ms", file=sys.stderr)
    finally:
        try:
            Path(out_path).unlink()
        except FileNotFoundError:
            pass

    return out


def diff_report(measured: dict, baseline: dict) -> tuple[str, bool]:
    rows = []
    has_regression = False
    for section in ("check_ms", "build_O0_ms"):
        for fixture, spec in baseline.get(section, {}).items():
            actual = measured[section].get(fixture)
            if actual is None:
                rows.append((f"{section}.{fixture}", spec["baseline_ms"], None,
                             None, spec["tolerance_pct"], "MISSING"))
                has_regression = True
                continue
            base = spec["baseline_ms"]
            tol = spec["tolerance_pct"]
            diff = 100.0 * (actual - base) / base if base else 0.0
            if diff > tol:
                status = "REGRESSION"
                has_regression = True
            elif diff < -tol:
                status = "improvement"
            else:
                status = "ok"
            rows.append((f"{section}.{fixture}", base, actual, diff, tol, status))

    lines = []
    lines.append(f"{'metric':<48} {'baseline':>10} {'actual':>10} {'diff':>9} {'tol':>6}  status")
    lines.append("-" * 100)
    for metric, base, actual, diff, tol, status in rows:
        if actual is None:
            lines.append(f"{metric:<48} {base:>10.1f} {'-':>10} {'-':>9} {tol:>5.1f}%  {status}")
        else:
            lines.append(
                f"{metric:<48} {base:>10.1f} {actual:>10.1f} "
                f"{diff:>+8.2f}% {tol:>5.1f}%  {status}"
            )
    regressed = sum(1 for r in rows if r[5] == "REGRESSION")
    improved = sum(1 for r in rows if r[5] == "improvement")
    stable = len(rows) - regressed - improved
    lines.append("")
    lines.append(f"summary: {len(rows)} metrics; "
                 f"{regressed} regressed, {improved} improved beyond tolerance, "
                 f"{stable} stable")
    return "\n".join(lines), has_regression


def write_baseline(measured: dict, baseline_path: Path) -> None:
    today = subprocess.check_output(["date", "+%Y-%m-%d"], text=True).strip()
    if baseline_path.exists():
        existing = json.loads(baseline_path.read_text())
        check_tols = {f: spec.get("tolerance_pct", 25)
                      for f, spec in existing.get("check_ms", {}).items()}
        build_tols = {f: spec.get("tolerance_pct", 20)
                      for f, spec in existing.get("build_O0_ms", {}).items()}
    else:
        check_tols = {}
        build_tols = {}

    new = {
        "version": 1,
        "captured": today,
        "captured_via": "bench/compile_check.py",
        "note": "Compile-time bench. `check_ms` is `ail check FILE`; `build_O0_ms` is `ail build --opt=-O0 FILE`. Wall-clock includes subprocess spawn (~5-10 ms on Linux) and, for build, the clang link step. Tolerances are tuned for noise on a quiet developer machine, not as the language correctness bar.",
        "check_ms": {},
        "build_O0_ms": {},
    }
    for fixture, ms in measured["check_ms"].items():
        new["check_ms"][fixture] = {
            "baseline_ms": round(ms, 2),
            "tolerance_pct": check_tols.get(fixture, 25),
        }
    for fixture, ms in measured["build_O0_ms"].items():
        new["build_O0_ms"][fixture] = {
            "baseline_ms": round(ms, 2),
            "tolerance_pct": build_tols.get(fixture, 20),
        }
    baseline_path.write_text(json.dumps(new, indent=2) + "\n")
    print(f">>> wrote new baseline to {baseline_path}", file=sys.stderr)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("-n", "--runs", type=int, default=5,
                    help="runs per fixture per op; min 2, default 5")
    ap.add_argument("--baseline", type=Path, default=DEFAULT_BASELINE)
    ap.add_argument("--update-baseline", action="store_true",
                    help="re-measure, then overwrite baseline_compile.json")
    args = ap.parse_args()
    if args.runs < 2:
        print("--runs must be >= 2", file=sys.stderr)
        return 2

    print(f">>> measuring {len(CORPUS)} fixtures, {args.runs} runs each "
          f"(check + build_O0)", file=sys.stderr)
    measured = measure(args.runs)

    if args.update_baseline:
        write_baseline(measured, args.baseline)
        return 0

    if not args.baseline.exists():
        print(f"no baseline at {args.baseline}; create one with --update-baseline",
              file=sys.stderr)
        return 2

    baseline = json.loads(args.baseline.read_text())
    report, has_regression = diff_report(measured, baseline)
    print(report)
    return 1 if has_regression else 0


if __name__ == "__main__":
    sys.exit(main())
