#!/usr/bin/env python3
"""
Per-operation latency harness for the bench_latency_{implicit,explicit}
fixtures.

The fixtures emit a stdout marker line every PRINT_K bench operations.
We spawn the binary with its stdout connected to a PTY so libc's
printf path is line-buffered (rather than block-buffered as it would
be on a pipe), record `time.monotonic_ns()` at every line read, and
report the inter-arrival distribution.

The headline numbers are median, p99, p99.9, max of the inter-arrival
gaps (in microseconds). Tail latency is the question; total wall time
is not.

We do NOT count the gap between READY (8888) and the first sample —
that gap includes the build_tree(19) cost and is not a steady-state
operation. We also drop the very last gap (DONE 9999) because the
program exit path can include final teardown that is not part of the
loop budget.

Usage:
    python3 bench/latency_harness.py <binary> [--label LABEL]
"""

from __future__ import annotations

import argparse
import os
import pty
import resource
import select
import sys
import time


def run_pty(binary: str, timeout_s: float = 600.0) -> tuple[list[int], int, int, int, float]:
    """Run `binary` with stdout on a PTY, return (timestamps_ns, lines, exit_status, max_rss_kb, wall_s)."""
    pid, fd = pty.fork()
    if pid == 0:
        # Child: replace argv[0] with the binary and exec.
        try:
            os.execvp(binary, [binary])
        except FileNotFoundError as e:
            sys.stderr.write(f"exec failed: {e}\n")
            os._exit(127)

    # Parent: read from the PTY master, timestamp every newline.
    timestamps: list[int] = []
    line_count = 0
    buf = bytearray()
    deadline = time.monotonic() + timeout_s
    t_start = time.monotonic()

    try:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                os.kill(pid, 9)
                raise TimeoutError(f"binary exceeded {timeout_s:.0f}s")
            r, _, _ = select.select([fd], [], [], min(remaining, 1.0))
            if not r:
                continue
            try:
                chunk = os.read(fd, 4096)
            except OSError:
                # PTY master read after child exit -> EIO on Linux.
                break
            if not chunk:
                break
            now = time.monotonic_ns()
            buf.extend(chunk)
            # Walk the buffer; record `now` for each newline encountered.
            # Multiple newlines in one read share the same monotonic_ns
            # — that's a known quantization artefact of bursty stdio
            # delivery, NOT a per-op timing artefact. We surface it in
            # the report.
            nl = buf.find(b"\n")
            while nl != -1:
                timestamps.append(now)
                line_count += 1
                del buf[: nl + 1]
                nl = buf.find(b"\n")
    finally:
        try:
            os.close(fd)
        except OSError:
            pass

    # Wait for child and capture rusage.
    _, status = os.waitpid(pid, 0)
    ru = resource.getrusage(resource.RUSAGE_CHILDREN)
    t_end = time.monotonic()
    return timestamps, line_count, status, ru.ru_maxrss, t_end - t_start


def percentile(xs: list[float], p: float) -> float:
    """Nearest-rank percentile (xs assumed sorted)."""
    if not xs:
        return float("nan")
    k = max(0, min(len(xs) - 1, int(round((p / 100.0) * (len(xs) - 1)))))
    return xs[k]


def summarise(timestamps_ns: list[int], lines: int) -> dict | None:
    """Compute the per-cell summary stats for one run. Returns None on
    failure (too few lines for a meaningful distribution)."""
    if lines < 5:
        return None
    gaps_ns: list[int] = []
    for i in range(1, len(timestamps_ns)):
        gaps_ns.append(timestamps_ns[i] - timestamps_ns[i - 1])
    # Trim the first gap (READY -> first chunk; includes the loop's
    # first chunk's full work) and the last gap (last chunk -> DONE;
    # includes program teardown). Everything between is steady-state.
    trimmed = gaps_ns[1:-1] if len(gaps_ns) >= 3 else gaps_ns
    gaps_us = sorted(g / 1000.0 for g in trimmed)
    if not gaps_us:
        return None
    return {
        "n": len(trimmed),
        "min": min(gaps_us),
        "median": percentile(gaps_us, 50),
        "mean": sum(gaps_us) / len(gaps_us),
        "p90": percentile(gaps_us, 90),
        "p99": percentile(gaps_us, 99),
        "p999": percentile(gaps_us, 99.9),
        "max": max(gaps_us),
        "zero_gaps": sum(1 for g in trimmed if g == 0),
    }


def cell_median(values: list[float]) -> tuple[float, float, float]:
    """Iter 18g.tidy.fu: median + min + max of a list of per-run values
    for one cell (e.g. p99 across N runs). Drops slowest-by-magnitude
    if N >= 4 to match `bench/run.sh`'s drop-slowest convention.
    Returns (median, min, max) — the orchestrator can read both the
    central tendency and the spread from one tuple.
    """
    if not values:
        return (float("nan"), float("nan"), float("nan"))
    s = sorted(values)
    if len(s) >= 4:
        # Drop the largest (slowest = highest latency) to absorb a
        # one-off STW pause or scheduler hiccup. Mirrors run.sh.
        s = s[:-1]
    n = len(s)
    if n % 2 == 1:
        med = s[n // 2]
    else:
        med = (s[n // 2 - 1] + s[n // 2]) / 2
    return (med, s[0], s[-1])


def report(
    label: str,
    runs: list[tuple[list[int], int, int, int, float]],
) -> None:
    """Iter 18g.tidy.fu: report one or many runs. With N=1 the output
    is byte-identical to the pre-stat-of-N format. With N>1 each cell
    shows median(min..max) across runs after dropping the slowest.
    """
    summaries: list[dict] = []
    last_status = 0
    last_rss = 0
    last_wall = 0.0
    last_lines = 0
    for timestamps, lines, status, rss_kb, wall_s in runs:
        last_status, last_rss, last_wall, last_lines = status, rss_kb, wall_s, lines
        s = summarise(timestamps, lines)
        if s is not None:
            summaries.append(s)

    if not summaries:
        print(f"[{label}] FAIL: no successful runs (lines={last_lines}, status={last_status})")
        return

    print(f"=== {label} ===")
    if len(summaries) == 1:
        # Single-run path: keep the pre-stat-of-N output exactly.
        s = summaries[0]
        print(f"  lines: {last_lines}  exit_status: {last_status}  wall: {last_wall:.3f}s  max_rss: {last_rss} KB")
        print(f"  inter-arrival gaps (steady-state, n={s['n']}, trimmed first+last):")
        print(f"    min:    {s['min']:10.1f} us")
        print(f"    median: {s['median']:10.1f} us")
        print(f"    mean:   {s['mean']:10.1f} us")
        print(f"    p90:    {s['p90']:10.1f} us")
        print(f"    p99:    {s['p99']:10.1f} us")
        print(f"    p99.9:  {s['p999']:10.1f} us")
        print(f"    max:    {s['max']:10.1f} us")
        print(f"    p99/median:   {s['p99'] / s['median'] if s['median'] else float('nan'):.2f}x")
        print(f"    max/median:   {s['max'] / s['median'] if s['median'] else float('nan'):.2f}x")
        print(f"  zero-ns gaps (read coalescing): {s['zero_gaps']} / {s['n']}")
        return

    # Multi-run path: median + min..max per cell across runs (drop
    # slowest if N >= 4, matching bench/run.sh).
    n_kept = len(summaries) - 1 if len(summaries) >= 4 else len(summaries)
    print(f"  runs: {len(summaries)}  kept: {n_kept} (drop slowest if N>=4)  last exit_status: {last_status}")
    print(f"  per-run distribution medians (steady-state, trimmed first+last):")

    def fmt_cell(key: str) -> str:
        med, lo, hi = cell_median([s[key] for s in summaries])
        return f"median={med:7.1f}  range=[{lo:7.1f}, {hi:7.1f}]"

    for key in ("min", "median", "mean", "p90", "p99", "p999", "max"):
        label_key = "p99.9" if key == "p999" else key
        print(f"    {label_key:7s} {fmt_cell(key)} us")

    # Derived ratios computed per-run, then aggregated.
    p99_med_ratios = [s["p99"] / s["median"] if s["median"] else float("nan") for s in summaries]
    max_med_ratios = [s["max"] / s["median"] if s["median"] else float("nan") for s in summaries]
    pmm_med, pmm_lo, pmm_hi = cell_median(p99_med_ratios)
    mmm_med, mmm_lo, mmm_hi = cell_median(max_med_ratios)
    print(f"    p99/median   median={pmm_med:.2f}x   range=[{pmm_lo:.2f}x, {pmm_hi:.2f}x]")
    print(f"    max/median   median={mmm_med:.2f}x   range=[{mmm_lo:.2f}x, {mmm_hi:.2f}x]")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("binary", help="path to the bench binary")
    ap.add_argument("--label", default=None, help="label for the report")
    ap.add_argument("--timeout", type=float, default=600.0)
    ap.add_argument(
        "--runs",
        type=int,
        default=1,
        help="number of timed runs (default 1). With N>=2 the report shows median + range per cell across runs; with N>=4 the slowest run is dropped before aggregation, matching bench/run.sh's drop-slowest convention.",
    )
    args = ap.parse_args()

    if args.runs < 1:
        print(f"--runs must be >= 1 (got {args.runs})", file=sys.stderr)
        return 2

    label = args.label or os.path.basename(args.binary)
    runs: list[tuple[list[int], int, int, int, float]] = []
    for _ in range(args.runs):
        runs.append(run_pty(args.binary, args.timeout))
    report(label, runs)
    return 0


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