#!/usr/bin/env python3
# Performance regression check.
#
# Runs bench/run.sh (or consumes its output from stdin / a file), parses
# the throughput table and the per-arm latency stanzas, and diffs every
# metric against bench/baseline.json. Exits 0 if every metric is within
# its per-metric tolerance; exits 1 with a per-metric diff table on any
# regression.
#
# Tolerances are one-sided: only metrics that got *worse* than baseline
# (slower wall-time, larger RSS, larger latency, larger ratio) trigger a
# regression. Improvements never fire — they show in the diff with a "+"
# sign and a note, so an intentional improvement can be ratified into
# the baseline with `--update-baseline`.
#
# Usage:
#     bench/check.py                       # spawn bench/run.sh -n 5
#     bench/check.py --from-file out.txt   # consume captured output
#     cat out.txt | bench/check.py --stdin
#     bench/check.py --update-baseline     # re-run, write new baseline
#     bench/check.py --baseline path.json  # alternate baseline location

from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_BASELINE = ROOT / "bench" / "baseline.json"
RUN_SH = ROOT / "bench" / "run.sh"


@dataclass
class Measurement:
    metric_id: str
    baseline: float
    tolerance_pct: float
    actual: float

    @property
    def diff_pct(self) -> float:
        if self.baseline == 0:
            return 0.0
        return 100.0 * (self.actual - self.baseline) / self.baseline

    @property
    def regressed(self) -> bool:
        return self.diff_pct > self.tolerance_pct


def parse_throughput_table(text: str) -> dict[str, dict[str, float]]:
    # The throughput table has a header row, a separator row, then one row
    # per workload. Columns are pipe-separated.
    out: dict[str, dict[str, float]] = {}
    in_table = False
    for line in text.splitlines():
        if line.startswith("workload") and "gc(s)" in line:
            in_table = True
            continue
        if in_table and line.startswith("---"):
            continue
        if in_table:
            if not line.strip() or "|" not in line:
                in_table = False
                continue
            cells = [c.strip() for c in line.split("|")]
            if len(cells) != 9:
                continue
            workload = cells[0]
            try:
                out[workload] = {
                    "gc_s":         float(cells[1]),
                    "bump_s":       float(cells[2]),
                    "rc_s":         float(cells[3]),
                    "gc_over_bump": float(cells[4].rstrip("x")),
                    "rc_over_bump": float(cells[5].rstrip("x")),
                    "gc_rss_kb":    float(cells[6]),
                    "bump_rss_kb":  float(cells[7]),
                    "rc_rss_kb":    float(cells[8]),
                }
            except ValueError:
                continue
    return out


# Latency arm header looks like: "=== implicit @ gc (Boehm-fair) ===".
# We map the leading prose label to the canonical arm key in baseline.json.
ARM_LABEL_TO_KEY = {
    "implicit @ gc": "implicit_at_gc",
    "explicit @ rc": "explicit_at_rc",
    "implicit @ rc": "implicit_at_rc",
}

# Inside each arm stanza, lines that report a metric look like:
#     "    median  median=   96.4  range=[..., ...] us"
# or:
#     "    p99/median   median=73.92x   range=[...]"
# We pick out the metric name and its "median=NN" value.
LINE_RE = re.compile(
    r"^\s+(?P<metric>[a-zA-Z0-9./_]+)\s+median=\s*(?P<value>[\d.]+)x?"
)
HEADER_RE = re.compile(r"^=== (?P<label>[^(]+?)\s*\(")

LATENCY_METRIC_RENAMES = {
    "median": "median_us",
    "p99": "p99_us",
    "p99.9": "p99_9_us",
    "max": "max_us",
    "p99/median": "p99_over_median",
}


def parse_latency_stanzas(text: str) -> dict[str, dict[str, float]]:
    out: dict[str, dict[str, float]] = {}
    current_key: str | None = None
    current: dict[str, float] = {}

    def flush() -> None:
        if current_key is not None and current:
            out.setdefault(current_key, {}).update(current)

    for line in text.splitlines():
        m = HEADER_RE.match(line)
        if m:
            flush()
            label = m.group("label").strip()
            current_key = ARM_LABEL_TO_KEY.get(label)
            current = {}
            continue
        if current_key is None:
            continue
        m = LINE_RE.match(line)
        if not m:
            continue
        raw_name = m.group("metric")
        canonical = LATENCY_METRIC_RENAMES.get(raw_name)
        if canonical is None:
            continue
        try:
            current[canonical] = float(m.group("value"))
        except ValueError:
            continue
    flush()
    return out


def run_bench() -> str:
    cmd = [str(RUN_SH), "-n", "5"]
    print(f">>> running: {' '.join(cmd)}", file=sys.stderr)
    proc = subprocess.run(
        cmd,
        cwd=str(ROOT),
        capture_output=True,
        text=True,
        check=False,
    )
    sys.stderr.write(proc.stderr)
    if proc.returncode != 0:
        print(f"bench/run.sh exited with code {proc.returncode}", file=sys.stderr)
        sys.exit(2)
    return proc.stdout


def collect_measurements(
    parsed_throughput: dict[str, dict[str, float]],
    parsed_latency: dict[str, dict[str, float]],
    baseline: dict,
) -> list[Measurement]:
    out: list[Measurement] = []
    missing: list[str] = []

    for workload, metrics in baseline.get("throughput", {}).items():
        actuals = parsed_throughput.get(workload)
        if actuals is None:
            missing.append(f"throughput.{workload}")
            continue
        for name, spec in metrics.items():
            actual = actuals.get(name)
            if actual is None:
                missing.append(f"throughput.{workload}.{name}")
                continue
            out.append(Measurement(
                metric_id=f"throughput.{workload}.{name}",
                baseline=spec["baseline"],
                tolerance_pct=spec["tolerance_pct"],
                actual=actual,
            ))

    for arm, metrics in baseline.get("latency", {}).items():
        actuals = parsed_latency.get(arm)
        if actuals is None:
            missing.append(f"latency.{arm}")
            continue
        for name, spec in metrics.items():
            actual = actuals.get(name)
            if actual is None:
                missing.append(f"latency.{arm}.{name}")
                continue
            out.append(Measurement(
                metric_id=f"latency.{arm}.{name}",
                baseline=spec["baseline"],
                tolerance_pct=spec["tolerance_pct"],
                actual=actual,
            ))

    if missing:
        print("error: parser did not find these baselined metrics in the bench output:",
              file=sys.stderr)
        for m in missing:
            print(f"  {m}", file=sys.stderr)
        print("Likely cause: bench/run.sh output format changed, or a fixture / latency arm was renamed. Update bench/check.py parsers or bench/baseline.json before claiming a regression.",
              file=sys.stderr)
        sys.exit(2)

    return out


def render_report(measurements: list[Measurement]) -> tuple[str, bool]:
    regressed = [m for m in measurements if m.regressed]
    improved  = [m for m in measurements if m.diff_pct < -m.tolerance_pct]

    lines = []
    lines.append(f"{'metric':<48} {'baseline':>12} {'actual':>12} {'diff':>9} {'tol':>6}  status")
    lines.append("-" * 100)
    for m in measurements:
        if m.regressed:
            status = "REGRESSION"
        elif m.diff_pct < -m.tolerance_pct:
            status = "improvement"
        else:
            status = "ok"
        lines.append(
            f"{m.metric_id:<48} {m.baseline:>12.3f} {m.actual:>12.3f} "
            f"{m.diff_pct:>+8.2f}% {m.tolerance_pct:>5.1f}%  {status}"
        )
    lines.append("")
    lines.append(f"summary: {len(measurements)} metrics; "
                 f"{len(regressed)} regressed, "
                 f"{len(improved)} improved beyond tolerance, "
                 f"{len(measurements) - len(regressed) - len(improved)} stable")
    return "\n".join(lines), bool(regressed)


def write_new_baseline(
    parsed_throughput: dict[str, dict[str, float]],
    parsed_latency: dict[str, dict[str, float]],
    baseline_path: Path,
    captured_via: str,
) -> None:
    today = subprocess.check_output(["date", "+%Y-%m-%d"], text=True).strip()
    existing = json.loads(baseline_path.read_text())
    new = {
        "version": existing.get("version", 1),
        "captured": today,
        "captured_via": captured_via,
        "note": existing.get("note", ""),
        "throughput": {},
        "latency": {},
    }
    for workload, metrics in existing.get("throughput", {}).items():
        actuals = parsed_throughput.get(workload, {})
        new["throughput"][workload] = {
            name: {
                "baseline": actuals.get(name, spec["baseline"]),
                "tolerance_pct": spec["tolerance_pct"],
            }
            for name, spec in metrics.items()
        }
    for arm, metrics in existing.get("latency", {}).items():
        actuals = parsed_latency.get(arm, {})
        new["latency"][arm] = {
            name: {
                "baseline": actuals.get(name, spec["baseline"]),
                "tolerance_pct": spec["tolerance_pct"],
            }
            for name, spec in metrics.items()
        }
    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__)
    src = ap.add_mutually_exclusive_group()
    src.add_argument("--from-file", type=Path, help="read bench output from this file")
    src.add_argument("--stdin", action="store_true", help="read bench output from stdin")
    ap.add_argument("--baseline", type=Path, default=DEFAULT_BASELINE)
    ap.add_argument("--update-baseline", action="store_true",
                    help="re-run, then overwrite baseline.json with the fresh numbers")
    args = ap.parse_args()

    if args.from_file:
        bench_output = args.from_file.read_text()
        captured_via = f"bench/run.sh (captured to {args.from_file})"
    elif args.stdin:
        bench_output = sys.stdin.read()
        captured_via = "bench/run.sh (piped to bench/check.py)"
    else:
        bench_output = run_bench()
        captured_via = "bench/run.sh -n 5"

    parsed_throughput = parse_throughput_table(bench_output)
    parsed_latency = parse_latency_stanzas(bench_output)
    baseline = json.loads(args.baseline.read_text())

    if args.update_baseline:
        write_new_baseline(parsed_throughput, parsed_latency, args.baseline, captured_via)
        return 0

    measurements = collect_measurements(parsed_throughput, parsed_latency, baseline)
    report, has_regression = render_report(measurements)
    print(report)
    return 1 if has_regression else 0


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