#!/usr/bin/env python3
"""Stage-by-stage timings for the hot-reload pipeline (1 ms polling)."""
import os, re, statistics, subprocess, sys, threading, time

ROOT = "/Users/oscar/.worktree/hot-reload/labs/hot-reload"
CLI = os.environ.get("KAN_HOT_CLI", "/Users/oscar/.worktree/hot-reload/target/debug/kan-hot-dev")
SOURCE = os.path.join(ROOT, "api-common", "src", "lib.rs")
ITER = int(sys.argv[1]) if len(sys.argv) > 1 else 8
EXTRA_ENV = dict(x.split("=", 1) for x in sys.argv[2:])

env = dict(os.environ, RUST_LOG="info", KAN_HOT_TIMING="1",
           KAN_HOT_EXAMPLE_INTERVAL_MS="50", **EXTRA_ENV)
process = subprocess.Popen([CLI, "run", "--bin", "server", "--lib", "api"],
                           cwd=ROOT, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                           stderr=subprocess.STDOUT, env=env, text=True, bufsize=1)
lines = []
lock = threading.Lock()


def reader():
    for line in process.stdout:
        with lock:
            lines.append((time.monotonic(), line))
        sys.stdout.write(line)
        sys.stdout.flush()


threading.Thread(target=reader, daemon=True).start()


def wait_for(pred, since=0.0, timeout=30.0):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        with lock:
            for stamp, line in reversed(lines):
                if stamp >= since and pred(line):
                    return stamp, line
        time.sleep(0.001)
    raise TimeoutError


wait_for(lambda l: "gen 1 hello" in l)
original = open(SOURCE).read()
stages = {k: [] for k in ("detect", "build", "link", "reload_cmd", "load", "marker", "total")}
try:
    for i in range(ITER):
        marker = f"stage-{os.getpid()}-{i}-{time.time_ns()}"
        t0 = time.monotonic()
        with open(SOURCE, "w") as f:
            f.write(original.replace("nested-v1", marker))
        t_detect, _ = wait_for(lambda l: "change detected" in l, since=t0)
        t_build, _ = wait_for(lambda l: "cargo hot build" in l, since=t_detect)
        t_link, _ = wait_for(lambda l: "thin link" in l, since=t_build)
        t_reload, _ = wait_for(lambda l: "kan-hot-dev: reloaded" in l, since=t_link)
        t_load, _ = wait_for(lambda l: "loaded generation" in l, since=t_reload)
        t_marker, _ = wait_for(lambda l: f'"nested":"{marker}"' in l, since=t_load)
        stages["detect"].append((t_detect - t0) * 1000)
        stages["build"].append((t_build - t_detect) * 1000)
        stages["link"].append((t_link - t_build) * 1000)
        stages["reload_cmd"].append((t_reload - t_link) * 1000)
        stages["load"].append((t_load - t_reload) * 1000)
        stages["marker"].append((t_marker - t_load) * 1000)
        stages["total"].append((t_marker - t0) * 1000)
        time.sleep(0.15)
finally:
    with open(SOURCE, "w") as f:
        f.write(original)
    time.sleep(0.5)
    try:
        process.stdin.write("q\n")
        process.stdin.flush()
    except BrokenPipeError:
        pass
    try:
        process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        process.kill()

for name, values in stages.items():
    print(f"{name:12s} median={statistics.median(values):6.1f} mean={statistics.mean(values):6.1f} "
          f"min={min(values):6.1f} max={max(values):6.1f}")