"""bench_daemon.py — daemon-mode vs per-spawn comparison.

Per-spawn pays ~200 ms CUDA context init every call. Daemon pays it
once at startup, then every subsequent call is just kernel + I/O.

For workloads doing many small calls (the natural shape of a
go-gpu primitive in a search loop), daemon mode is the only mode
that wins."""

import hashlib
import os
import secrets
import subprocess
import sys
import time

BIN = sys.argv[1] if len(sys.argv) > 1 else "./shake256-fanout"
OUT_BYTES = 32


def write_portal(path, inputs_hex, out_bytes):
    with open(path, "w") as fh:
        fh.write("(cuda-shake-fanout\n")
        fh.write(f"  (output-bytes {out_bytes})\n")
        fh.write("  (inputs\n")
        for h in inputs_hex:
            fh.write(f'    "{h}"\n')
        fh.write("))\n")


def main():
    print(f"binary: {BIN}\n")
    N_INPUTS = 100   # per-call (small workload)
    N_CALLS  = 10    # number of fan-out calls

    # generate N_CALLS distinct input sets
    portals = []
    for i in range(N_CALLS):
        hex_in = [secrets.token_hex(32) for _ in range(N_INPUTS)]
        in_p = f"/tmp/cf-d-in-{i}.portal"
        write_portal(in_p, hex_in, OUT_BYTES)
        portals.append((in_p, f"/tmp/cf-d-out-{i}.portal"))

    # ── per-spawn (control) ────────────────────────────────────────
    t0 = time.monotonic()
    for in_p, out_p in portals:
        subprocess.run([BIN, in_p, out_p], check=True)
    spawn_t = time.monotonic() - t0

    # ── daemon mode ────────────────────────────────────────────────
    t0 = time.monotonic()
    proc = subprocess.Popen([BIN, "--daemon"],
                            stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
    ready = proc.stdout.readline().strip()
    assert ready == "ready", f"daemon didn't ready: {ready!r}"
    init_t = time.monotonic() - t0

    t0 = time.monotonic()
    for in_p, out_p in portals:
        proc.stdin.write(f"process {in_p} {out_p}\n")
        proc.stdin.flush()
        line = proc.stdout.readline().strip()
        assert line.startswith("done"), f"daemon error: {line!r}"
    daemon_calls_t = time.monotonic() - t0

    proc.stdin.write("quit\n"); proc.stdin.flush()
    proc.stdout.readline(); proc.wait()
    daemon_total = init_t + daemon_calls_t

    # ── host hashlib (reference) ───────────────────────────────────
    inputs_hex = [secrets.token_hex(32) for _ in range(N_INPUTS)]
    t0 = time.monotonic()
    for _ in range(N_CALLS):
        [hashlib.shake_256(bytes.fromhex(h)).hexdigest(OUT_BYTES) for h in inputs_hex]
    host_t = time.monotonic() - t0

    # cleanup
    for in_p, out_p in portals:
        try: os.unlink(in_p)
        except: pass
        try: os.unlink(out_p)
        except: pass

    # ── report ─────────────────────────────────────────────────────
    print(f"workload: {N_CALLS} calls × {N_INPUTS} inputs × 32 bytes each\n")
    print(f"  host hashlib loop  : {host_t*1e3:>8.1f} ms total ({host_t*1e3/N_CALLS:.2f} ms/call)")
    print(f"  per-spawn fanout   : {spawn_t*1e3:>8.1f} ms total ({spawn_t*1e3/N_CALLS:.2f} ms/call)")
    print(f"  daemon-mode init   : {init_t*1e3:>8.1f} ms  (one-time)")
    print(f"  daemon-mode calls  : {daemon_calls_t*1e3:>8.1f} ms total "
          f"({daemon_calls_t*1e3/N_CALLS:.2f} ms/call)")
    print(f"  daemon end-to-end  : {daemon_total*1e3:>8.1f} ms (init + calls)")
    print()
    print(f"  daemon/per-spawn   : {daemon_total/spawn_t:.3f}× ({(1 - daemon_total/spawn_t)*100:+.1f}% wall time)")
    print(f"  daemon/host        : {daemon_total/host_t:.2f}× "
          f"({'GPU wins' if daemon_total < host_t else 'host wins'})")
    print()
    if N_CALLS >= 2:
        avg_per_call = daemon_calls_t / N_CALLS
        spawn_per_call = spawn_t / N_CALLS
        print(f"  per-call comparison (excluding daemon init):")
        print(f"    per-spawn: {spawn_per_call*1e3:.2f} ms each (cuda ctx init each)")
        print(f"    daemon   : {avg_per_call*1e3:.2f} ms each (warm context)")
        print(f"    daemon is {spawn_per_call/avg_per_call:.1f}× faster per call")


if __name__ == "__main__":
    main()
