"""bench_tiers.py — round-trip latency, Python vs C vs asm worker.

Spawns one worker per tier (each backed by the same warm daemon
binary), times N round-trips at small + medium inputs, reports
median + p99 + Σ.

Usage:
  python3 bench_tiers.py [--small N] [--medium N]
  python3 bench_tiers.py --skip-asm   ; if asm host-side helpers not wired

What we're measuring:
  - WIRE orchestration cost (S-expression parse, portal write,
    pipe-to-daemon, daemon response, portal read, S-expression format)
  - NOT the CUDA kernel — that's identical across tiers since each
    tier shells to the same shake256-fanout --daemon binary

So the differences come from: TCP recv/send buffering, S-expression
reader/printer speed, file I/O wrapping, pipe write/flush handling.
"""

import os
import secrets
import socket
import struct
import subprocess
import sys
import time

ROOT = os.path.dirname(os.path.abspath(__file__))
LUMBDA_ROOT = os.path.abspath(os.path.join(ROOT, "..", ".."))

TIERS = {
    "python": [
        "python3", "-u", os.path.join(LUMBDA_ROOT, "lumbda.py"),
        "launch.lsp", "--port",
    ],
    "c": [
        os.path.join(LUMBDA_ROOT, "c", "lumbda"),
        "launch.lsp", "--port",
    ],
    # asm tier launch uses a separate file that pre-defines *argv*
    # (asm doesn't auto-bind it). Port is fixed to 8320 — asm tier
    # also doesn't parse --port, so each asm bench uses its own
    # binary launch.
    #
    # Port mnemonic — 8320 = BEND:
    #   8 ~= B (implied infinity B flattened; bake a cake; baby & me)
    #   3 ~= E (backward)
    #   2 ~= N (pivoted 90 degrees)
    #   0 ~= D (flattened)
    "asm": [
        os.path.join(LUMBDA_ROOT, "asm", "lumbda-gc"),
        "launch-asm.lsp",
    ],
}


def wire_send(sock, payload):
    hdr = f"{len(payload):08d}".encode("ascii")
    sock.sendall(hdr + payload.encode("utf-8"))


def wire_recv(sock):
    hdr = recv_exact(sock, 8)
    if hdr is None:
        return None
    plen = int(hdr.decode())
    return recv_exact(sock, plen).decode("utf-8")


def recv_exact(sock, n):
    buf = b""
    while len(buf) < n:
        c = sock.recv(n - len(buf))
        if not c:
            return None
        buf += c
    return buf


def spawn_worker(tier, port):
    cmd = TIERS[tier].copy()
    if tier != "asm":
        cmd.append(str(port))
    proc = subprocess.Popen(
        cmd, cwd=ROOT,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
    )
    # Wait for "listening on port" or timeout
    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        line = proc.stdout.readline()
        if not line:
            break
        if "listening on port" in line:
            return proc
    proc.kill()
    raise RuntimeError(f"{tier}: worker didn't ready ({line!r})")


def stop_worker(proc):
    proc.terminate()
    try:
        proc.wait(timeout=3)
    except subprocess.TimeoutExpired:
        proc.kill()


def time_calls(port, payload, n, binary=False):
    """N round-trips, return list of per-call seconds."""
    times = []
    for _ in range(n):
        s = socket.socket()
        s.connect(("127.0.0.1", port))
        t0 = time.monotonic()
        if binary:
            hdr = f"{len(payload):08d}".encode("ascii")
            s.sendall(hdr + payload)
            rh = recv_exact_sock(s, 8)
            n_resp = int(rh.decode())
            resp = recv_exact_sock(s, n_resp)
        else:
            wire_send(s, payload)
            resp = wire_recv(s)
        elapsed = time.monotonic() - t0
        s.close()
        if binary:
            if not resp or resp[:4] != b"BSHR":
                raise RuntimeError(f"bad binary resp: {resp[:80]!r}")
        else:
            if resp is None or "(ok" not in resp:
                raise RuntimeError(f"bad response: {resp!r}")
        times.append(elapsed)
    return times


def recv_exact_sock(sock, n):
    buf = b""
    while len(buf) < n:
        c = sock.recv(n - len(buf))
        if not c:
            return None
        buf += c
    return buf


def make_payload(n_inputs, input_bytes):
    inputs = [secrets.token_hex(input_bytes) for _ in range(n_inputs)]
    parts = ['(cuda-shake-fanout (quote (']
    for h in inputs:
        parts.append(f'"{h}" ')
    parts.append(f")) 32)")
    return "".join(parts)


def make_payload_binary(n_inputs, input_bytes, out_bytes=32):
    """Binary wire payload: 'BSHK' + u32 out_bytes + u32 n + n x (u32 len + bytes).
    Sent as raw bytes — no hex encoding, no S-expression."""
    blobs = [secrets.token_bytes(input_bytes) for _ in range(n_inputs)]
    parts = [b"BSHK", struct.pack("<II", out_bytes, n_inputs)]
    for b in blobs:
        parts.append(struct.pack("<I", len(b)))
        parts.append(b)
    return b"".join(parts)


def stats(times):
    s = sorted(times)
    n = len(s)
    return {
        "n": n,
        "min_ms": s[0] * 1e3,
        "med_ms": s[n // 2] * 1e3,
        "p99_ms": s[min(n - 1, int(n * 0.99))] * 1e3,
        "max_ms": s[-1] * 1e3,
        "tot_s": sum(s),
    }


def bench_tier(tier, port, configs, binary=False):
    print(f"\n=== {tier} tier (port {port}{' BINARY' if binary else ''}) ===")
    proc = spawn_worker(tier, port)
    # Warmup so CUDA context init + daemon spawn doesn't skew first call
    time.sleep(0.5)
    warmup_payload = (make_payload_binary(3, 4) if binary
                      else make_payload(3, 4))
    for _ in range(5):
        time_calls(port, warmup_payload, 1, binary=binary)
    results = {}
    for label, (n_calls, n_inputs, input_bytes) in configs.items():
        p = (make_payload_binary(n_inputs, input_bytes) if binary
             else make_payload(n_inputs, input_bytes))
        try:
            t = time_calls(port, p, n_calls, binary=binary)
        except Exception as e:
            # Workers can cliff at large payloads; mark this cell &
            # respawn so subsequent (potentially smaller) cells in
            # the same row still get measured.
            print(f"  {label:20}  CLIFF  ({e})")
            results[label] = None
            try:
                stop_worker(proc)
            except Exception:
                pass
            try:
                proc = spawn_worker(tier, port)
                time.sleep(0.5)
                for _ in range(2):
                    time_calls(port, warmup_payload, 1, binary=binary)
            except Exception as e2:
                print(f"  (respawn failed: {e2})")
                return results
            continue
        st = stats(t)
        results[label] = st
        print(f"  {label:20} n={st['n']:>3}  med={st['med_ms']:>6.2f} ms  "
              f"p99={st['p99_ms']:>6.2f} ms  min={st['min_ms']:>5.2f} ms")
    stop_worker(proc)
    return results


def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--skip-asm", action="store_true",
                    help="skip asm tier (if host-side helpers not wired)")
    ap.add_argument("--small-n", type=int, default=100)
    ap.add_argument("--medium-n", type=int, default=30)
    ap.add_argument("--binary", action="store_true",
                    help="use binary wire mode (BSHK protocol)")
    args = ap.parse_args()

    # Write a launch.lsp the worker process will load
    with open(os.path.join(ROOT, "launch.lsp"), "w") as f:
        f.write('(load "wire.lsp")\n(load "gpu-worker.lsp")\n(main)\n')
    # asm tier doesn't auto-bind *argv* OR ship write-to-string;
    # pre-define the former and shim the latter via wts-shim.
    with open(os.path.join(ROOT, "launch-asm.lsp"), "w") as f:
        f.write('(define *argv* (quote ()))\n'
                '(load "write-to-string-shim.lsp")\n'
                '(load "wire.lsp")\n(load "gpu-worker.lsp")\n(main)\n')

    if args.binary:
        # Binary mode pushes through CLIFFs since the worker bypasses
        # S-expression parsing entirely. Workload sizes shifted higher.
        configs = {
            "small   (100 × 16 B)":   (args.small_n,   100,        16),
            "medium  (1000 × 16 B)":  (args.medium_n,  1000,       16),
            "med     (10k × 16 B)":   (10,             10_000,     16),
            "huge    (100k × 16 B)":  (5,              100_000,    16),
            "huge    (1M × 16 B)":    (3,              1_000_000,  16),
        }
    else:
        configs = {
            "small  (3 × 16 B)":     (args.small_n,   3,        16),
            "small  (100 × 16 B)":   (args.small_n,   100,      16),
            "medium (1000 × 16 B)":  (args.medium_n,  1000,     16),
            "med    (10k × 16 B)":   (3,              10_000,   16),
            "huge   (50k × 16 B)":   (3,              50_000,   16),
            "huge   (100k × 16 B)":  (2,              100_000,  16),
        }

    tiers = ["python", "c"]
    if not args.skip_asm:
        tiers.append("asm")

    all_results = {}
    for i, tier in enumerate(tiers):
        # asm tier hard-codes port 8320 (BEND, no --port parsing); run
        # it last so other tiers get unique ports.
        port = 8320 if tier == "asm" else (9090 + i + 1)
        try:
            all_results[tier] = bench_tier(tier, port, configs, binary=args.binary)
        except Exception as e:
            print(f"  {tier} FAILED: {e}")
            all_results[tier] = None

    # Comparative table
    print("\n=== summary (median ms) ===")
    print(f"{'workload':<22}", end="")
    for tier in tiers:
        print(f"  {tier:>10}", end="")
    print()
    for label in configs:
        print(f"{label:<22}", end="")
        for tier in tiers:
            tr = all_results.get(tier)
            if tr is None:
                print(f"  {'-':>10}", end="")
            else:
                cell = tr.get(label)
                if cell is None:
                    print(f"  {'CLIFF':>10}", end="")
                else:
                    print(f"  {cell['med_ms']:>9.2f}", end="")
        print()


if __name__ == "__main__":
    main()
