#!/bin/bash
# rpc-chain-bench.sh — benchmark S-expression RPC chains across impls.
#
# Six layers of safety (see CLAUDE.md #Asm memory discipline):
#   1. set -e            — fail-fast
#   2. ulimit -v 524288  — 512 MB virt cap, kernel-enforced
#   3. trap on EXIT/INT/TERM — pkill any leftover server
#   4. timeout 30        — wall-clock cap on every server spawn
#   5. explicit kill + wait at end of each bench
#   6. pgrep verify before moving on — hard stop if a straggler is left
#
# Measures:
#   A. Direct:  Python client → asm backend              (1 hop)
#   B. 1-relay: Python client → C relay → asm backend    (2 hops)
#   C. 2-relay: Python → Py relay → C relay → asm        (3 hops)
#
# Per-hop cost = (B - A) / N requests = C relay overhead per request.

set -e
cd "$(dirname "$0")/.."
ulimit -v 524288

SPAWNED=()
cleanup() {
    for pid in "${SPAWNED[@]}"; do kill -9 "$pid" 2>/dev/null; done
    sleep 0.2
    pkill -9 -u "$USER" -f 'lumbda.*(rpc|repl)-(server|relay|chain-bench)\.lsp' 2>/dev/null || true
}
trap cleanup EXIT INT TERM

start_bg() {
    # $1=cmd. Spawns with timeout; returns PID on stdout.
    local cmd="$1"
    eval "timeout 30 $cmd >/tmp/srv_$$.log 2>&1 &"
    local pid=$!
    SPAWNED+=("$pid")
    echo "$pid"
}

wait_port() {
    # $1=port. Polls until something accepts on that port (or timeout).
    local port="$1" i
    for i in $(seq 1 40); do
        if timeout 0.2 bash -c "</dev/tcp/127.0.0.1/$port" 2>/dev/null; then return 0; fi
        sleep 0.1
    done
    return 1
}

stop_bg() {
    local pid="$1"
    kill "$pid" 2>/dev/null || true
    for _ in 1 2 3 4 5; do kill -0 "$pid" 2>/dev/null || break; sleep 0.1; done
    kill -9 "$pid" 2>/dev/null || true
    wait "$pid" 2>/dev/null || true
    # Remove from SPAWNED array
    local new=(); for p in "${SPAWNED[@]}"; do [ "$p" != "$pid" ] && new+=("$p"); done
    SPAWNED=("${new[@]}")
    sleep 0.3
}

verify_clean() {
    local strays
    strays=$(pgrep -u "$USER" -f 'lumbda.*(rpc|repl)-(server|relay|chain-bench)\.lsp' 2>/dev/null || true)
    if [ -n "$strays" ]; then
        echo "!! STRAYS: $strays" >&2
        pkill -9 -u "$USER" -f 'lumbda.*(rpc|repl)-(server|relay|chain-bench)\.lsp' 2>/dev/null || true
        return 1
    fi
}

PY="python3 lumbda.py --fast"
C="./c/lumbda"
ASM="./asm/lumbda"

# Always use Python client to drive (identical across benches).
# We sed the hardcoded port in-place because asm's top-level
# `define` inside a loaded file overrides any wrapper-set value.
client_bench() {
    local port="$1"
    sed "s/^(define \*target-port\*.*/(define *target-port* $port)/" \
        examples/rpc-chain-bench.lsp > /tmp/bench-drv.lsp
    timeout 20 $PY /tmp/bench-drv.lsp 2>&1 | grep -E "^(ok|elapsed|rps)" | head -3
}

# Same, but using any client impl
client_bench_with() {
    local port="$1" client="$2"
    sed "s/^(define \*target-port\*.*/(define *target-port* $port)/" \
        examples/rpc-chain-bench.lsp > /tmp/bench-drv.lsp
    if [[ "$client" == "asm" ]]; then
        timeout 20 $ASM < /tmp/bench-drv.lsp 2>&1 | grep -E "^(ok|elapsed|rps)" | head -3
    elif [[ "$client" == "C" ]]; then
        timeout 20 $C /tmp/bench-drv.lsp 2>&1 | grep -E "^(ok|elapsed|rps)" | head -3
    else
        timeout 20 $PY /tmp/bench-drv.lsp 2>&1 | grep -E "^(ok|elapsed|rps)" | head -3
    fi
}

echo "═══════════════════════════════════════════════════════════════"
echo "RPC chain benchmark — 200 ping requests"
echo "═══════════════════════════════════════════════════════════════"
echo

# ── A. Direct: Py client → asm backend on 9080 ──
echo "── (A) Python client → asm backend (direct, 1 hop) ──"
BPID=$(start_bg "$ASM < examples/rpc-server.lsp")
wait_port 9080 || { echo "backend failed to start"; exit 1; }
client_bench 9080
stop_bg "$BPID"
verify_clean
echo

# ── B. 1 relay: Py client → C relay (9082) → asm backend (9080) ──
echo "── (B) Python client → C relay → asm backend (2 hops) ──"
BPID=$(start_bg "$ASM < examples/rpc-server.lsp")
wait_port 9080 || { echo "backend failed"; exit 1; }

# Relay wrapper generator: rewrites the hardcoded ports. Avoids the
# pre-define-then-load trick because asm's `define` at top level
# redefines (shadows) a pre-existing binding.
make_relay() {
    local listen="$1" backend="$2" out="$3"
    sed -e "s/^(define \*listen-port\*.*/(define *listen-port* $listen)/" \
        -e "s/^(define \*backend-port\*.*/(define *backend-port* $backend)/" \
        examples/rpc-relay.lsp > "$out"
}

make_relay 9082 9080 /tmp/relay-cr.lsp
RPID=$(start_bg "$C /tmp/relay-cr.lsp")
wait_port 9082 || { echo "relay failed"; exit 1; }
client_bench 9082
stop_bg "$RPID"
stop_bg "$BPID"
verify_clean
echo

# ── C. 2 relays: Py client → Py relay → C relay → asm backend ──
echo "── (C) Python client → Py relay → C relay → asm (3 hops) ──"
BPID=$(start_bg "$ASM < examples/rpc-server.lsp")
wait_port 9080 || { echo "backend failed"; exit 1; }

make_relay 9082 9080 /tmp/relay-cr.lsp
RPID1=$(start_bg "$C /tmp/relay-cr.lsp")
wait_port 9082 || { echo "relay 1 failed"; exit 1; }

make_relay 9083 9082 /tmp/relay-pr.lsp
RPID2=$(start_bg "$PY /tmp/relay-pr.lsp")
wait_port 9083 || { echo "relay 2 failed"; exit 1; }
client_bench 9083
stop_bg "$RPID2"
stop_bg "$RPID1"
stop_bg "$BPID"
verify_clean
echo

# ── D. Reverse chain: asm client → Python relay → C relay → asm backend ──
echo "── (D) asm client → Py relay → C relay → asm backend (3 hops) ──"
BPID=$(start_bg "$ASM < examples/rpc-server.lsp")
wait_port 9080 || { echo "backend failed"; exit 1; }

make_relay 9082 9080 /tmp/relay-cr.lsp
RPID1=$(start_bg "$C /tmp/relay-cr.lsp")
wait_port 9082 || { echo "relay 1 failed"; exit 1; }

make_relay 9083 9082 /tmp/relay-pr.lsp
RPID2=$(start_bg "$PY /tmp/relay-pr.lsp")
wait_port 9083 || { echo "relay 2 failed"; exit 1; }

client_bench_with 9083 asm
stop_bg "$RPID2"
stop_bg "$RPID1"
stop_bg "$BPID"
verify_clean
echo

rm -f /tmp/bench-drv.lsp /tmp/relay-cr.lsp /tmp/relay-pr.lsp /tmp/srv_$$.log
echo "═══════════════════════════════════════════════════════════════"
echo "Done. Safety net fired zero times (no strays)."