#!/usr/bin/env python3
"""Interleaved A/B benchmark. Variant syntax: name[@ENV=VAL;ENV2=VAL2]"""
import os, re, statistics, subprocess, sys

BENCH = "/private/var/folders/ft/fkt3bwqd01g7f54wvyz9wgjc0000gn/T/opencode/p4/bench_ab.py"
BIN = "/private/var/folders/ft/fkt3bwqd01g7f54wvyz9wgjc0000gn/T/opencode/p4/bin/"
LAB = "/Users/oscar/.worktree/hot-reload/labs/hot-reload"
ROUNDS = int(sys.argv[1]) if len(sys.argv) > 1 else 2
ITERS = sys.argv[2] if len(sys.argv) > 2 else "6"
SPECS = sys.argv[3].split(",") if len(sys.argv) > 3 else ["base", "new"]


def parse(spec):
    name, _, rest = spec.partition("@")
    extra = dict(pair.split("=", 1) for pair in rest.split(";") if pair)
    return spec, name, extra


def run(name, extra):
    env = dict(os.environ, KAN_HOT_CLI=BIN + name, KAN_HOT_LAB_ROOT=LAB, **extra)
    out = subprocess.run([sys.executable, BENCH, ITERS], env=env, cwd=LAB,
                         capture_output=True, text=True).stdout
    result = {}
    for key in ("end-to-end", "cargo hot build", "thin link", "dlopen"):
        m = re.search(rf"^{key}[^:]*: n=\d+ mean=([\d.]+)ms median=([\d.]+)ms min=([\d.]+)ms max=([\d.]+)ms",
                      out, re.M)
        if m:
            result[key] = tuple(float(x) for x in m.groups())
    return result


specs = [parse(s) for s in SPECS]
# Warm up: the first run of a binary may pay a Cargo rebuild.
for spec, name, extra in specs:
    run(name, extra)

results = {spec: [] for spec, _, _ in specs}
for round_ in range(ROUNDS):
    for spec, name, extra in specs:
        r = run(name, extra)
        results[spec].append(r)
        print(f"round {round_} {spec}: " + " | ".join(f"{k} med={v[1]:.1f}" for k, v in r.items()),
              flush=True)

print()
for spec, runs in results.items():
    print(f"== {spec}")
    for key in ("end-to-end", "cargo hot build", "thin link", "dlopen"):
        values = [r[key] for r in runs if key in r]
        if values:
            print(f"  {key:18s} " + "  ".join(f"{v[1]:.1f}" for v in values)
                  + f"   mean-of-medians={statistics.mean(v[1] for v in values):.1f}")