#!/usr/bin/env python3
"""Time individual replayed rustc invocations with flag variations."""
import json, os, subprocess, sys, time, statistics, shutil

REC = "/Users/oscar/.worktree/hot-reload/labs/hot-reload/target/kan-hot/debug/rustc"
ROOT = "/Users/oscar/.worktree/hot-reload/labs/hot-reload"
SRC_COMMON = os.path.join(ROOT, "api-common/src/lib.rs")
SRC_API = os.path.join(ROOT, "api/src/lib.rs")


def load(name):
    return json.load(open(os.path.join(REC, name + ".json")))


def run(inv, extra=(), mutate=None, timeout=120):
    argv = list(inv["argv"]) + list(extra)
    env = {k: v for k, v in inv["env"]}
    env.pop("RUSTC_BOOTSTRAP", None)
    started = time.monotonic()
    p = subprocess.run(argv, cwd=inv["cwd"], env=env,
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    elapsed = (time.monotonic() - started) * 1000
    if p.returncode != 0:
        raise SystemExit(f"failed: {p.stderr.decode()[:3000]}\nargv={argv}")
    return elapsed


def mutate(path, marker):
    original = open(path).read()
    if "nested-v1" in original:
        open(path, "w").write(original.replace("nested-v1", marker))
    else:
        open(path, "w").write(original + f"\n// {marker}\n")


def bench(label, name, variants, n=7):
    inv = load(name)
    path = SRC_COMMON if name == "api_common" else SRC_API
    results = {}
    for vlabel, extra, cleanup in variants:
        if cleanup:
            for d in cleanup:
                shutil.rmtree(d, ignore_errors=True)
        times = []
        for i in range(n):
            mutate(path, f"bench-{time.time_ns()}")
            times.append(run(inv, extra))
        results[vlabel] = times
        print(f"{label:14s} {vlabel:34s} median={statistics.median(times):6.1f} "
              f"mean={statistics.mean(times):6.1f} min={min(times):6.1f} max={max(times):6.1f}")
    return results


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 7
    inc_common = os.path.join(ROOT, "target/bench-incremental-common")
    inc_api = os.path.join(ROOT, "target/bench-incremental-api")
    variants_common = [
        ("baseline cgu16", (), ()),
        ("cgu1", ("-C", "codegen-units=1"), ()),
        ("cgu4", ("-C", "codegen-units=4"), ()),
        ("cgu8", ("-C", "codegen-units=8"), ()),
        ("cgu32", ("-C", "codegen-units=32"), ()),
        ("incremental", ("-C", f"incremental={inc_common}"), (inc_common,)),
        ("incremental cgu1", ("-C", f"incremental={inc_common}", "-C", "codegen-units=1"), (inc_common,)),
        ("no dep-info", ("--emit=metadata,link",), ()),
        ("incremental+no depinfo", ("-C", f"incremental={inc_common}", "--emit=metadata,link"), (inc_common,)),
    ]
    variants_api = [
        ("baseline cgu16", (), ()),
        ("cgu1", ("-C", "codegen-units=1"), ()),
        ("cgu8", ("-C", "codegen-units=8"), ()),
        ("cgu32", ("-C", "codegen-units=32"), ()),
        ("incremental", ("-C", f"incremental={inc_api}"), (inc_api,)),
        ("incremental cgu1", ("-C", f"incremental={inc_api}", "-C", "codegen-units=1"), (inc_api,)),
        ("no dep-info", ("--emit=metadata,link",), ()),
        ("incremental+no depinfo", ("-C", f"incremental={inc_api}", "--emit=metadata,link"), (inc_api,)),
    ]
    if which in ("all", "common"):
        bench("api_common", "api_common", variants_common, n)
    if which in ("all", "api"):
        bench("api", "api", variants_api, n)