#!/usr/bin/env bash
# pp-test.sh — Qwen3.8-27B pipeline-parallel (llama.cpp RPC) test orchestrator.
#
# Runs Qwen3.8-27B UD-Q8_K_XL split across Dankclaw (head, Strix iGPU) and
# Dankserver (worker, RTX 3090) over the 2.5 GbE LAN. Both boxes run the same
# llama.cpp build (b10433 / 9b05354ec); the worker build lives in
# /home/dank/llama-cpp-rpc on Dankserver.
#
# BENCHMARKING POLICY (AGENTS.md): this test STOPS both production services
# and restores them when done. Do not run it while users are mid-stream.
#
# Usage (from any box with ssh access to both targets, repo checked out):
#   ./scripts/pp-test.sh preflight
#   ./scripts/pp-test.sh start                # stop prod, start PP head+worker
#   ./scripts/pp-test.sh verify               # show the layer split + VRAM
#   ./scripts/pp-test.sh bench [prompt_k] [out_tokens] [reps]
#   ./scripts/pp-test.sh stop-pp              # stop PP head+worker only
#   ./scripts/pp-test.sh restore              # start production services
#   ./scripts/pp-test.sh all                  # start + verify + bench + restore
#
# Sudo: set FRANK_SUDO_PW for headless runs (stop/start of system services).
#   FRANK_SUDO_PW=... ./scripts/pp-test.sh start
set -euo pipefail

CLAW="dank@dankclaw"
SERVER="dank@dankserver"
CLAW_REPO="/home/dank/vllmScripts"
SERVER_REPO="/home/dank/vLLM"
HEAD_PORT=7770
HEAD_LOG="/tmp/pp-head.log"
WORKER_LOG="/tmp/pp-worker.log"
RPC_PORT=50052

SUDO="sudo"
if [ -n "${FRANK_SUDO_PW:-}" ]; then
    SUDO="printf '%s\n' '$FRANK_SUDO_PW' | sudo -S"
fi

log()  { echo "[pp-test] $*"; }
die()  { echo "[pp-test] FAIL: $*" >&2; exit 1; }

ssh_claw()   { ssh -o BatchMode=yes -o ConnectTimeout=8 "$CLAW"   "$@"; }
ssh_server() { ssh -o BatchMode=yes -o ConnectTimeout=8 "$SERVER" "$@"; }

# ---------------------------------------------------------------------------
# preflight
# ---------------------------------------------------------------------------
do_preflight() {
    log "preflight: build identity"
    local h_v w_v
    h_v=$(ssh_claw '/home/dank/llama.cpp/build/bin/llama-server --version 2>&1 | head -1')
    w_v=$(ssh_server '/home/dank/llama-cpp-rpc/build/bin/llama-server --version 2>&1 | head -1')
    log "  head:   $h_v"
    log "  worker: $w_v"
    [[ "$h_v" == *9b05354ec* && "$w_v" == *9b05354ec* ]] || die "build skew (need identical commit 9b05354ec on both sides)"

    log "preflight: binaries"
    ssh_claw   "[ -x $CLAW_REPO/scripts/serve-pp-head.sh ]"   || die "missing $CLAW_REPO/scripts/serve-pp-head.sh on claw (git pull?)"
    ssh_server "[ -x $SERVER_REPO/scripts/serve-pp-worker.sh ]" || die "missing $SERVER_REPO/scripts/serve-pp-worker.sh on server (git pull?)"
    ssh_server "[ -x /home/dank/llama-cpp-rpc/build/bin/ggml-rpc-server ]" || die "missing worker RPC binary on server"
    ssh_claw   "[ -f /home/dank/models/Qwen3.8-27B-UD-Q8_K_XL.gguf ]" || die "missing model on claw"

    log "preflight: network (claw -> server:50052 listener check needs the worker up; testing route + link)"
    ssh_claw "ip -4 addr show eno1 | grep -q 192.168.0.143" || die "unexpected claw address"
    ssh_server "ip -4 addr show eno1 | grep -q 192.168.0.235" || die "unexpected server address"
    ssh_claw "ping -c 2 -W 2 192.168.0.235 >/dev/null" || die "claw cannot reach server"
    local rtt
    rtt=$(ssh_claw "ping -c 5 -W 2 192.168.0.235 | tail -1" | grep -oE '[0-9.]+/[0-9.]+ ms' | head -1 || true)
    log "  rtt: ${rtt:-unknown}"

    log "preflight: production services (health endpoint = no sudo needed)"
    local h1 h2
    h1=$(ssh_claw   "curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:7770/health 2>/dev/null || true")
    h2=$(ssh_server "curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:7770/health 2>/dev/null || true")
    log "  frank-default-large:  ${h1:-no-response}"
    log "  frank-default-small:  ${h2:-no-response}"
    log "preflight OK"
}

# ---------------------------------------------------------------------------
# start
# ---------------------------------------------------------------------------
do_start() {
    log "stopping production services (both boxes)"
    ssh_claw   "$SUDO systemctl stop frank-default-large"  || die "could not stop frank-default-large"
    ssh_server "$SUDO systemctl stop frank-default-small"  || die "could not stop frank-default-small"
    sleep 5

    log "dropping page caches on claw (reliable model load)"
    ssh_claw "$SUDO sh -c 'sync; echo 3 > /proc/sys/vm/drop_caches'" || log "WARN: cache drop failed"

    log "starting RPC worker on server"
    ssh_server "cd $SERVER_REPO && nohup bash scripts/serve-pp-worker.sh > $WORKER_LOG 2>&1 & echo worker_pid=\$!"
    for i in $(seq 1 30); do
        if ssh_server "ss -tln | grep -q :$RPC_PORT"; then
            log "worker listening on :$RPC_PORT"
            break
        fi
        [ "$i" = "30" ] && { ssh_server "tail -20 $WORKER_LOG"; die "worker did not bind :$RPC_PORT"; }
        sleep 1
    done

    log "starting PP head on claw (model load + weight stream: allow ~10 min)"
    local head_pid
    head_pid=$(ssh_claw "cd $CLAW_REPO && nohup bash scripts/serve-pp-head.sh > $HEAD_LOG 2>&1 & echo \$!")
    log "head pid on claw: $head_pid"
    for i in $(seq 1 360); do
        local st
        st=$(ssh_claw "curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:$HEAD_PORT/health 2>/dev/null || true")
        if [ "$st" = "200" ]; then
            log "head healthy after ~$((i * 5))s"
            return 0
        fi
        if ! ssh_claw "kill -0 $head_pid 2>/dev/null"; then
            ssh_claw "tail -30 $HEAD_LOG"
            die "head process died during load"
        fi
        [ $((i % 12)) = 0 ] && log "  still loading... ($((i * 5))s)"
        sleep 5
    done
    ssh_claw "tail -30 $HEAD_LOG"
    die "head did not become healthy in 30 min"
}

# ---------------------------------------------------------------------------
# verify
# ---------------------------------------------------------------------------
do_verify() {
    log "health + device split (from head log)"
    ssh_claw "curl -s http://127.0.0.1:$HEAD_PORT/health; echo"
    ssh_claw "grep -aE 'using device|assigned to device|MiB free' $HEAD_LOG | head -20 || true"
    log "worker VRAM on server:"
    ssh_server "nvidia-smi --query-gpu=memory.used --format=csv,noheader"
    log "worker log tail:"
    ssh_server "tail -5 $WORKER_LOG"
}

# ---------------------------------------------------------------------------
# bench
# ---------------------------------------------------------------------------
do_bench() {
    local PROMPT_K="${1:-4}"
    local OUT_TOKENS="${2:-256}"
    local REPS="${3:-2}"
    log "bench: prompt ~${PROMPT_K}k tokens, out=${OUT_TOKENS}, reps=${REPS} (on claw, via :$HEAD_PORT)"
    ssh_claw "python3 $CLAW_REPO/scripts/pp-bench.py $PROMPT_K $OUT_TOKENS $REPS $HEAD_PORT $HEAD_LOG"
}

# ---------------------------------------------------------------------------
# stop / restore
# ---------------------------------------------------------------------------
do_stop_pp() {
    log "stopping PP head (claw) and worker (server)"
    ssh_claw   "pkill -f 'serve-pp-head.sh' 2>/dev/null; pkill -f 'llama-server.*frank-pp-large' 2>/dev/null; true"
    ssh_server "pkill -x ggml-rpc-server 2>/dev/null; true"
    sleep 3
    ssh_server "nvidia-smi --query-gpu=memory.used --format=csv,noheader"
}

do_restore() {
    # Warm-cache loads crash fast on the claw (AGENTS.md); cold loads survive.
    log "dropping page caches on claw before production load"
    ssh_claw "$SUDO sh -c 'sync; echo 3 > /proc/sys/vm/drop_caches'" || log "WARN: cache drop failed"
    log "starting production services"
    ssh_server "$SUDO systemctl start frank-default-small" || die "could not start frank-default-small"
    ssh_claw   "$SUDO systemctl start frank-default-large" || die "could not start frank-default-large"
    log "waiting for health (model load takes minutes; up to ~17 min cold on claw)"
    for i in $(seq 1 300); do
        local a b
        a=$(ssh_claw   "curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:7770/health 2>/dev/null || true")
        b=$(ssh_server "curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:7770/health 2>/dev/null || true")
        if [ "$a" = "200" ] && [ "$b" = "200" ]; then
            log "both production services healthy after ~$((i * 4))s"
            return 0
        fi
        [ $((i % 15)) = 0 ] && log "  claw=$a server=$b ... ($((i * 4))s)"
        sleep 4
    done
    die "production services did not both become healthy"
}

# ---------------------------------------------------------------------------
case "${1:-}" in
    preflight) do_preflight ;;
    start)     do_start ;;
    verify)    do_verify ;;
    bench)     shift; do_bench "$@" ;;
    stop-pp)   do_stop_pp ;;
    restore)   do_restore ;;
    all)
        log "full run: any failure auto-restores production services"
        do_preflight
        trap 'do_restore || true' EXIT
        do_start
        do_verify
        do_bench
        do_restore
        trap - EXIT
        ;;
    *)
        echo "Usage: $0 {preflight|start|verify|bench [prompt_k] [out] [reps]|stop-pp|restore|all}"
        exit 1
        ;;
esac