#!/usr/bin/env bash
# frank-archon-arms.sh — run a SEQUENCE of frank-archon bench arms and keep
# every arm's evidence. Runs on the operator box; each arm is one full window
# of ./autoresearch.sh (stop+mask -> boot the arm as the service -> measure ->
# restore production and verify), so the fleet is only ever down for one arm at
# a time and the production unit is back up between arms.
#
# An arm is the repo's own frank-archon.env plus an OVERLAY of engine knobs
# appended after it (later assignment wins when the serve script sources it),
# staged through FRANK_ARCHON_AR_ENV_FILE so the checkout is never mutated and
# each arm gets its own fingerprint (the memo can never answer one arm with
# another arm's numbers).
#
# Usage:  scripts/autoresearch/frank-archon-arms.sh [arm ...]
#         scripts/autoresearch/frank-archon-arms.sh --list
# Output, under /tmp/frank-archon-arms-<UTC stamp>/ (override with ARMS_DIR):
#
#   <arm>.out        full autoresearch output (metrics + engine evidence)
#   <arm>.env        the exact staged env
#   <arm>.rc         window exit code
#   summary.tsv      arm, rc, decode_tps, decode_tps_t0.7, kv_tokens, boot_s
#
# Exit 0 when every arm produced metrics; 1 if any arm failed (a failed boot is
# a RESULT — the window still restored production, and the rc says which arm).
set -uo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd -- "$SCRIPT_DIR/../.." && pwd)"
BASE_ENV="$REPO/scripts/frank-archon.env"
ARMS_DIR="${ARMS_DIR:-/tmp/frank-archon-arms-$(date -u +%Y%m%dT%H%M%SZ)}"

# --- arms ------------------------------------------------------------------
# baseline: the shipped shape E, measured with THIS harness so the comparison is
# matched (the ~29.6 t/s figure in the docs is an older shape, and no shape-E
# decode number exists yet).
OVER_baseline=()
# B: MTP3 with the vendor's probabilistic draft sampling, a capture ladder whose
# step is num_speculative_tokens + 1 = 4, and the sampling override. Precise
# memory stays ON — this arm is simultaneously the MTP test and the test of
# whether MTP's drafter workspace and the fitted 5.25 GiB overhead can coexist.
SPEC3='{"method":"mtp","num_speculative_tokens":3,"draft_sample_method":"probabilistic"}'
LADDER3='{"cudagraph_capture_sizes":[4,8,12],"max_cudagraph_capture_size":12,"inductor_compile_config":{"combo_kernels":false,"benchmark_combo_kernel":false}}'
# The value must arrive in the staged env file SINGLE-QUOTED: sourcing an env
# file performs quote removal, so an unquoted {"temperature":0.7} would reach
# the engine as {temperature:0.7}.
GEN07='{"temperature":0.7}'
OVER_B=(
    "FRANK_ARCHON_SPEC='$SPEC3'"
    "FRANK_ARCHON_COMPCONFIG='$LADDER3'"
    'FRANK_ARCHON_GEN_CONFIG='"'$GEN07'"
)
# B2: identical to B with the default (7.60 GiB) overhead, i.e. precise memory
# OFF. It is the fallback if B dies at graph capture, and it prices the arena
# that precise memory buys (7,522 slots -> ~6,496).
OVER_B2=(
    "FRANK_ARCHON_SPEC='$SPEC3'"
    "FRANK_ARCHON_COMPCONFIG='$LADDER3'"
    'FRANK_ARCHON_GEN_CONFIG='"'$GEN07'"
    'FRANK_ARCHON_PRECISE_MEMORY='
)
# C: the vendor's concurrency shape — SEQS 16 with the [4..32] ladder and util
# 0.95. This is also the direct test of the theory that pinning the capture
# ladder unlocks SEQS > 3 (the SEQS=2 failure was query_start_loc specialised to
# the constant 3, which is NOT in the default ladder; 4 is).
LADDER16='{"cudagraph_capture_sizes":[4,8,12,16,20,24,28,32],"max_cudagraph_capture_size":32,"inductor_compile_config":{"combo_kernels":false,"benchmark_combo_kernel":false}}'
OVER_C=(
    "FRANK_ARCHON_SPEC='$SPEC3'"
    "FRANK_ARCHON_COMPCONFIG='$LADDER16'"
    'FRANK_ARCHON_GEN_CONFIG='"'$GEN07'"
    'FRANK_ARCHON_SEQS=16'
    'FRANK_ARCHON_MNBT=4096'
    'FRANK_ARCHON_UTIL=0.95'
)

# D/E: the CONTEXT arms. Measured 2026-10-02 (arms run 20261002T233855Z): the
# KV pool is a hard byte request divided by a PER-TOKEN PRICE, and MTP raises
# that price because the drafter keeps its own KV per position:
#   no spec: 3.15e9 / 398,658 tokens =  7,901 B/token  (1.99x of 200,704)
#   MTP3:    3.15e9 / 313,002 tokens = 10,064 B/token  (1.56x)
# So arm C lost a fifth of the context to the drafter, not to util. Frank's
# contract is >=401,400 tokens over >=2 seqs = 2.00x of 200,704, which at
# 10,064 B/token costs 4.04e9 -> ask 4.05e9 (+0.83 GiB/rank).
# D is arm B with EXACTLY ONE variable changed (the KV bytes), so its result
# says whether the 4k-depth collapse (B 23.09 vs baseline 33.99) was KV or the
# batch shape.
KV401K='FRANK_ARCHON_KV_BYTES=4050000000'
OVER_D=(
    "FRANK_ARCHON_SPEC='$SPEC3'"
    "FRANK_ARCHON_COMPCONFIG='$LADDER3'"
    'FRANK_ARCHON_GEN_CONFIG='"'$GEN07'"
    "$KV401K"
)
# E: D with C's max_num_batched_tokens (4096 — C is the only arm that did NOT
# collapse at 4k depth: 33.15 where B gave 23.09) and util 0.99, the value
# Frank asked about. 0.99 is affordable on paper — production at util 0.98
# measures 30.37/34.21 GiB per card, 3.58 GiB physically unused — but it leaves
# only 0.32 GiB above the allowance while the CUDA-graph cost is an ANALYTIC
# estimate (C's was 1.47 GiB over 24 graphs), so it is measured, not assumed.
OVER_E=(
    "FRANK_ARCHON_SPEC='$SPEC3'"
    "FRANK_ARCHON_COMPCONFIG='$LADDER3'"
    'FRANK_ARCHON_GEN_CONFIG='"'$GEN07'"
    "$KV401K"
    'FRANK_ARCHON_MNBT=4096'
    'FRANK_ARCHON_UTIL=0.99'
)

arm_overlay() {
    case "$1" in
        baseline) ;;                       # the shipped env, no overlay
        B)  printf '%s\n' "${OVER_B[@]}" ;;
        B2) printf '%s\n' "${OVER_B2[@]}" ;;
        C)  printf '%s\n' "${OVER_C[@]}" ;;
        D)  printf '%s\n' "${OVER_D[@]}" ;;
        E)  printf '%s\n' "${OVER_E[@]}" ;;
        *) return 1 ;;
    esac
}
if [ "${1:-}" = "--list" ]; then
    printf 'baseline B B2 C D E\n'
    exit 0
fi

[ -f "$BASE_ENV" ] || { echo "[arms] missing $BASE_ENV" >&2; exit 1; }
if [ "$#" -gt 0 ]; then ARMS=("$@"); else ARMS=(baseline B B2 C); fi
mkdir -p "$ARMS_DIR"
echo "[arms] dir $ARMS_DIR  arms: ${ARMS[*]}" >&2

: >"$ARMS_DIR/summary.tsv"
printf 'arm\trc\tdecode_tps\tdecode_tps_t0.7\tkv_tokens\tboot_s\n' >>"$ARMS_DIR/summary.tsv"
overall=0

for arm in "${ARMS[@]}"; do
    if ! arm_overlay "$arm" >"$ARMS_DIR/$arm.overlay"; then
        echo "[arms] unknown arm: $arm" >&2
        overall=1
        continue
    fi
    { cat "$BASE_ENV"; echo; echo "# --- arm $arm overlay ($(date -u +%FT%TZ)) ---"
      arm_overlay "$arm"; } >"$ARMS_DIR/$arm.env"
    # Gate the staged env BEFORE burning a ~20-minute window: every knob that
    # holds JSON must still be one valid JSON object after the file is SOURCED
    # (that is the step that silently eats unquoted braces).
    if ! bash -c 'set -a; . "$1" >/dev/null; set +a
        for k in FRANK_ARCHON_SPEC FRANK_ARCHON_COMPCONFIG FRANK_ARCHON_GEN_CONFIG; do
            v="${!k:-}"
            [ -n "$v" ] || continue
            case "$v" in
                "{"*) python3 -c "import json,sys; json.loads(sys.argv[1])" "$v" \
                        || { echo "[arms] $k is not valid JSON: $v" >&2; exit 1; } ;;
                *) echo "[arms] $k is not a JSON object: $v" >&2; exit 1 ;;
            esac
        done' _ "$ARMS_DIR/$arm.env"; then
        echo "[arms] $arm: staged env rejected, window not opened" >&2
        echo 99 >"$ARMS_DIR/$arm.rc"
        overall=1
        continue
    fi

    echo "[arms] === $arm: opening window $(date -u +%H:%M:%S)Z ===" >&2
    FRANK_ARCHON_AR_ENV_FILE="$ARMS_DIR/$arm.env" \
        "$REPO/autoresearch.sh" >"$ARMS_DIR/$arm.out" 2>&1
    rc=$?
    echo "$rc" >"$ARMS_DIR/$arm.rc"
    echo "[arms] === $arm: rc=$rc $(date -u +%H:%M:%S)Z ===" >&2
    grep -E '^METRIC ' "$ARMS_DIR/$arm.out" | sed 's/^METRIC //' >"$ARMS_DIR/$arm.metrics" || true
    metric() { local v; v="$(sed -n "s/^$1=//p" "$ARMS_DIR/$arm.metrics" | tail -1)"; printf '%s' "${v:-$2}"; }
    printf '%s\t%s\t%s\t%s\t%s\t%s\n' "$arm" "$rc" \
        "$(metric decode_tps FAIL)" "$(metric decode_tps_t0.7 FAIL)" \
        "$(metric kv_tokens ?)" "$(metric boot_seconds ?)" >>"$ARMS_DIR/summary.tsv"
    [ "$rc" -ne 0 ] && overall=1
done

echo "[arms] summary:" >&2
cat "$ARMS_DIR/summary.tsv" >&2
exit "$overall"