#!/usr/bin/env python3
"""Measure edit -> new code serving latency for the hot-reload example.

Runs kan-hot-dev, edits api-common with unique markers, and waits until the
server prints the new marker. Also extracts the internal build/link/load
breakdown from KAN_HOT_TIMING and kan-hot's load log.

Usage: python3 labs/hot-reload/bench_reload.py [iterations]
"""
import os
import re
import statistics
import subprocess
import sys
import threading
import time

ROOT = os.environ.get("KAN_HOT_LAB_ROOT", os.path.dirname(os.path.abspath(__file__)))
CLI = os.environ["KAN_HOT_CLI"]
SOURCE = os.path.join(ROOT, "api-common", "src", "lib.rs")


def main():
    iterations = int(sys.argv[1]) if len(sys.argv) > 1 else 5
    if not os.path.exists(CLI):
        sys.exit(f"build kan-hot-dev first: {CLI}")

    env = dict(
        os.environ,
        RUST_LOG="info",
        KAN_HOT_TIMING="1",
        KAN_HOT_EXAMPLE_INTERVAL_MS="50",
    )
    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(predicate, timeout=30.0):
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            with lock:
                for stamp, line in reversed(lines):
                    if predicate(line):
                        return stamp, line
            time.sleep(0.005)
        raise TimeoutError("timed out waiting for output")

    wait_for(lambda line: "gen 1 hello" in line)
    original = open(SOURCE).read()
    samples = []
    builds = []
    links = []
    loads = []

    try:
        for iteration in range(iterations):
            marker = f"bench-{os.getpid()}-{iteration}-{time.time_ns()}"
            with open(SOURCE, "w") as handle:
                handle.write(original.replace("nested-v1", marker))
            start = time.monotonic()
            wait_for(
                lambda line, marker=marker: f'"nested":"{marker}"' in line
            )
            samples.append(time.monotonic() - start)

            with lock:
                recent = [line for _, line in lines[-40:]]
            for line in recent:
                match = re.search(r"cargo hot build ([0-9.]+)(ms|s|µs)", line)
                if match:
                    builds.append(_seconds(match))
                match = re.search(r"thin link ([0-9.]+)(ms|s|µs)", line)
                if match:
                    links.append(_seconds(match))
                match = re.search(r"loaded generation \d+ in ([0-9.]+)(ms|s|µs)", line)
                if match:
                    loads.append(_seconds(match))
            time.sleep(0.15)
    finally:
        with open(SOURCE, "w") as handle:
            handle.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()

    def report(name, values):
        if not values:
            return
        print(
            f"\n{name}: n={len(values)} mean={statistics.mean(values):.1f}ms "
            f"median={statistics.median(values):.1f}ms "
            f"min={min(values):.1f}ms max={max(values):.1f}ms"
        )

    report("end-to-end edit -> serving", [value * 1000 for value in samples])
    report("cargo hot build", [value * 1000 for value in builds])
    report("thin link", [value * 1000 for value in links])
    report("dlopen", [value * 1000 for value in loads])


def _seconds(match):
    value = float(match.group(1))
    unit = match.group(2)
    return value * (1e-6 if unit == "µs" else 1e-3 if unit == "ms" else 1.0)


if __name__ == "__main__":
    main()