#!/usr/bin/env python3
"""Corpus wall-clock overhead harness for nash (docs/NASH.md §11, M2/M3 gate: <3%).

Replays the corpus under bash and nash in alternating rounds and compares total
wall-clock per shell. Only the shell subprocess is timed (fixture seeding and
filesystem snapshots are outside the clock). Rounds alternate shell order so
cache/thermal drift cancels; the reported figure uses the median round total.

Usage: overhead.py [--nash PATH] [--bash PATH] [--corpus PATH] [--rounds N]
                   [--observe-spool DIR] [--gate PCT] [--json PATH]

--observe-spool enables nash observation (spool transport) so the measured
configuration is the deployed one; the env is set identically for bash, where
it is inert.
"""

import argparse
import json
import os
import shutil
import statistics
import subprocess
import sys
import tempfile
import time

from replay import FIXTURES, TIMEOUT_S, seed


def run_timed(shell, cmd, extra_env):
    parent = tempfile.mkdtemp(prefix="nash-overhead-")
    workdir = os.path.join(parent, "workspace")
    os.makedirs(workdir)
    seed(workdir)
    env = {
        "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
        "HOME": workdir,
        "LC_ALL": "C",
        "LANG": "C",
        "TERM": "dumb",
        "SHELL": shell,
    }
    if extra_env:
        env.update(extra_env)
    start = time.monotonic()
    try:
        subprocess.run(
            [shell, "-c", cmd],
            cwd=workdir,
            env=env,
            capture_output=True,
            timeout=TIMEOUT_S,
        )
    except subprocess.TimeoutExpired:
        pass
    elapsed = time.monotonic() - start
    shutil.rmtree(parent, ignore_errors=True)
    return elapsed


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--nash", default=os.environ.get("NASH_BIN", "nash"))
    ap.add_argument("--bash", default="/bin/bash")
    ap.add_argument(
        "--corpus",
        default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "corpus.jsonl"),
    )
    ap.add_argument("--rounds", type=int, default=5)
    ap.add_argument("--observe-spool", help="enable nash observation, spooling to this dir")
    ap.add_argument("--gate", type=float, default=3.0, help="max overhead percent")
    ap.add_argument("--json", help="also write results to this path")
    args = ap.parse_args()

    observe_env = None
    if args.observe_spool:
        os.makedirs(args.observe_spool, exist_ok=True)
        observe_env = {
            "NUCLEIC_SHELL_SPOOL": args.observe_spool,
            "NUCLEIC_SESSION_ID": "corpus-overhead",
        }

    with open(args.corpus) as f:
        cmds = [json.loads(line)["cmd"] for line in f if line.strip()]

    # Warm-up: one untimed pass per shell (page cache, binary load).
    for shell in (args.bash, args.nash):
        for cmd in cmds[:10]:
            run_timed(shell, cmd, observe_env)

    totals = {"bash": [], "nash": []}
    for round_no in range(args.rounds):
        order = [("bash", args.bash), ("nash", args.nash)]
        if round_no % 2:
            order.reverse()
        for name, shell in order:
            total = sum(run_timed(shell, cmd, observe_env) for cmd in cmds)
            totals[name].append(total)
            print(f"round {round_no + 1} {name}: {total:.3f}s", flush=True)

    bash_med = statistics.median(totals["bash"])
    nash_med = statistics.median(totals["nash"])
    overhead = 100.0 * (nash_med / bash_med - 1.0)

    print(f"\ncorpus: {len(cmds)} cmds x {args.rounds} rounds")
    print(f"bash median: {bash_med:.3f}s   nash median: {nash_med:.3f}s")
    print(f"overhead: {overhead:+.2f}%   (gate: <{args.gate:.1f}%)")

    if args.json:
        with open(args.json, "w") as f:
            json.dump(
                {
                    "cmds": len(cmds),
                    "rounds": args.rounds,
                    "totals": totals,
                    "bash_median_s": bash_med,
                    "nash_median_s": nash_med,
                    "overhead_percent": overhead,
                    "gate_percent": args.gate,
                    "observed": bool(args.observe_spool),
                },
                f,
                indent=1,
            )

    gate = overhead < args.gate
    print(f"M3 overhead gate (<{args.gate:.1f}%): {'PASS' if gate else 'FAIL'}")
    return 0 if gate else 1


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