#!/usr/bin/env bash
# frank-archon-window.sh — one full bench window for frank-archon, run ON
# DANKSERVER as the deploy user (invoked over ssh by ./autoresearch.sh).
#
# Discipline (AGENTS.md Benchmarking Policy): the production unit is
# STOPPED + MASKED first, the candidate config boots AS THE SERVICE itself
# (same serve script, same container name, same port 7770 — no scratch
# ports, no co-tenant GPU), the in-container measure runs, and the window
# ALWAYS ends by restoring the production unit and verifying its health.
# The EXIT trap makes the restore unconditional: a boot failure, a measure
# failure, or a dead ssh session all end with production back up.
#
# Staging layout (created by the orchestrator, ephemeral — never a config
# source; the repo checkout is):
#   /home/dank/.autoresearch-frank-archon/
#     serve.sh             staged candidate serve script
#     frank-archon.env    staged candidate env (the candidate knobs)
#     measure.py           the in-container measure client
#     hash                 candidate fingerprint
#     run.json             result (metrics + boot + kv + fingerprint)
#     serve.out/measure.out/fail.log
#     last/                memo: run.json + hash of the last FULLY
#                          successful window (bench AND verified restore)
#
# This script logs to stderr and exits 0 only when the measure produced a
# usable result (the restore happens in the trap regardless of the bench
# outcome; a failed restore forces exit 1).
set -uo pipefail

STAGE="${FRANK_ARCHON_AR_STAGE:-/home/dank/.autoresearch-frank-archon}"
UNIT="${FRANK_ARCHON_AR_UNIT:-frank-archon}"
CTR="${FRANK_ARCHON_AR_CTR:-frank-archon-qwen}"
HEALTH_URL="http://127.0.0.1:7770/health"
MODELS_URL="http://127.0.0.1:7770/v1/models"
BOOT_TIMEOUT="${FRANK_ARCHON_AR_BOOT_TIMEOUT:-1200}"
RESTORE_TIMEOUT="${FRANK_ARCHON_AR_RESTORE_TIMEOUT:-1200}"
MEASURE_TIMEOUT="${FRANK_ARCHON_AR_MEASURE_TIMEOUT:-1800}"
SERVED="$(sed -n 's/^FRANK_ARCHON_SERVED_NAME=//p' "$STAGE/frank-archon.env" 2>/dev/null | head -1)"
SERVED="${SERVED:-frank-archon}"

RESTORED=0

log() { printf '[window %s] %s\n' "$(date '+%H:%M:%S')" "$*" >&2; }

container_state() { podman inspect -f '{{.State.Status}}' "$CTR" 2>/dev/null || echo absent; }

wait_port_free() {
    local i
    for i in $(seq 1 90); do
        [ "$(container_state)" != "running" ] && return 0
        sleep 2
    done
    return 1
}

wait_health() { # $1 = timeout seconds
    local t0 code
    t0=$(date +%s)
    while :; do
        code="$(curl -s -o /dev/null -w '%{http_code}' --max-time 10 "$HEALTH_URL" 2>/dev/null || echo 000)"
        [ "$code" = "200" ] && return 0
        [ "$(( $(date +%s) - t0 ))" -ge "$1" ] && return 1
        sleep 2
    done
}

capture_bench_failure() {
    log "capturing bench-boot failure evidence"
    {
        echo "=== podman logs (tail 400) ==="
        podman logs --tail 400 "$CTR" 2>&1 || true
        echo "=== staged serve.out (tail 60) ==="
        tail -60 "$STAGE/serve.out" 2>/dev/null || true
        echo "=== engine log (tail 60) ==="
        tail -60 /home/dank/logs/frank_archon.log 2>/dev/null || true
    } >"$STAGE/fail.log" 2>&1 || true
    grep -E 'OutOfMemory|Tried to allocate|Engine core initialization failed|Error|error' \
        "$STAGE/fail.log" 2>/dev/null | tail -25 >&2 || true
}

restore() {
    [ "$RESTORED" = 1 ] && return 0
    RESTORED=1
    log "restoring production unit $UNIT"
    if [ -f "$STAGE/serve.pid" ]; then
        kill "$(cat "$STAGE/serve.pid" 2>/dev/null)" 2>/dev/null || true
    fi
    podman stop -t 120 "$CTR" >/dev/null 2>&1 || true
    wait_port_free || log "WARNING: bench container still draining after 180s"
    sudo systemctl unmask "$UNIT" 2>/dev/null || true
    if ! sudo systemctl start "$UNIT"; then
        log "ERROR: systemctl start $UNIT failed — operator intervention required"
        return 1
    fi
    if wait_health "$RESTORE_TIMEOUT"; then
        if curl -s --max-time 10 "$MODELS_URL" 2>/dev/null | grep -q "$SERVED"; then
            log "production restored and verified ($SERVED answering on :7770)"
            # Memo promotion lives HERE: a cached result is only ever the
            # product of a fully successful window (bench AND verified
            # restore), never of a bare bench success.
            cp "$STAGE/run.json" "$STAGE/last/" 2>/dev/null || true
            cp "$STAGE/hash" "$STAGE/last/" 2>/dev/null || true
        else
            log "ERROR: health 200 but /v1/models does not list $SERVED"
            return 1
        fi
    else
        log "ERROR: production not healthy within ${RESTORE_TIMEOUT}s — unit left to systemd (Restart=always)"
        return 1
    fi
    return 0
}
trap 'restore; rc=$?; [ "$rc" -ne 0 ] && exit 1' EXIT

main() {
    local f
    for f in serve.sh frank-archon.env measure.py hash; do
        [ -f "$STAGE/$f" ] || { log "staged file missing: $STAGE/$f"; exit 1; }
    done
    mkdir -p "$STAGE/last"
    rm -f "$STAGE/run.json" "$STAGE/fail.log" "$STAGE/measure.out" "$STAGE/evidence.log"

    local t_start t0 i kvline kv_tokens kv_conc measure_json boot_seconds
    t_start="$(date +%s)"

    log "stopping and masking production unit $UNIT"
    sudo systemctl stop "$UNIT" 2>/dev/null || true
    # Explicit stop is not auto-restarted even with Restart=always; the mask
    # is belt-and-braces against a stray start during the window.
    sudo systemctl mask "$UNIT" 2>/dev/null || \
        log "WARNING: mask failed — relying on explicit-stop semantics"
    wait_port_free || log "note: production container still draining"

    t0="$(date +%s)"
    log "starting bench instance from staged config (fingerprint $(tr -d '[:space:]' <"$STAGE/hash"))"
    FRANK_ARCHON_ENV_FILE="$STAGE/frank-archon.env" \
        setsid nohup bash "$STAGE/serve.sh" >"$STAGE/serve.out" 2>&1 </dev/null &
    echo "$!" >"$STAGE/serve.pid"

    for i in $(seq 1 45); do
        [ "$(container_state)" = "running" ] && break
        sleep 2
    done
    if [ "$(container_state)" != "running" ]; then
        log "bench container never reached running state"
        tail -40 "$STAGE/serve.out" >&2 || true
        exit 1
    fi

    if wait_health "$BOOT_TIMEOUT"; then
        boot_seconds=$(( $(date +%s) - t0 ))
        log "bench healthy after ${boot_seconds}s"
    else
        if [ "$(container_state)" = "running" ]; then
            log "bench boot timed out after ${BOOT_TIMEOUT}s (still loading)"
        else
            log "bench container died during boot"
            capture_bench_failure
        fi
        exit 1
    fi

    log "running measure in container (timeout ${MEASURE_TIMEOUT}s)"
    if ! podman cp "$STAGE/measure.py" "$CTR":/tmp/ar-measure.py; then
        log "podman cp of measure client failed"
        exit 1
    fi
    if ! timeout "$MEASURE_TIMEOUT" podman exec "$CTR" python3 /tmp/ar-measure.py \
            >"$STAGE/measure.out" 2>&1; then
        log "measure failed or timed out"
        tail -40 "$STAGE/measure.out" >&2 || true
        exit 1
    fi

    # Capacity facts for the run (context-size experiments change these).
    kvline="$(podman logs "$CTR" 2>&1 | grep -E 'GPU KV cache size' | tail -1 || true)"
    kv_tokens="$(printf '%s' "$kvline" | sed -E 's/.*GPU KV cache size: ([0-9,]+) tokens.*/\1/' | tr -d ',' || true)"
    kv_conc="$(printf '%s' "$kvline" | sed -E 's/.*per request: ([0-9.]+x).*/\1/' || true)"

    # Human-facing evidence, in a sidecar file because it is multi-line prose,
    # not metrics: the expert-plan lines (the ONLY place the arena / runtime
    # overhead / KV-reserve numbers exist — they are never in /metrics), the
    # compile path (cold build vs cache hit), the graph-capture sizes, the
    # speculative-decode counters, and any OOM signature. Without this an arm
    # that loses decode speed is unexplainable afterwards.
    {
        echo "=== engine plan / capacity / spec evidence ==="
        podman logs "$CTR" 2>&1 | grep -E \
'KV cache size|Maximum concurrency|layers split|expert cache|runtime overhead|non-default args|r4d all-reduce|Using \[|cudagraph|CUDA graph|Capturing|graph for compile range|torch.compile took|Directly load the compiled|speculative|Speculative|acceptance|Acceptance|offload|precise|OutOfMemory|Tried to allocate|Engine core initialization failed' \
            | tail -150 || true
        echo "=== /metrics counters (spec decode, caches, success) ==="
        curl -s --max-time 15 http://127.0.0.1:7770/metrics 2>/dev/null \
            | grep -E '^vllm:(spec_decode|gpu_cache|prefix_cache|request_success|generation_tokens|iteration)' \
            | head -80 || true
    } >"$STAGE/evidence.log" 2>&1 || true
    log "engine evidence -> $STAGE/evidence.log ($(wc -l <"$STAGE/evidence.log") lines)"

    measure_json="$(grep -E '^\{' "$STAGE/measure.out" | tail -1 || true)"
    if [ -z "$measure_json" ]; then
        log "measure produced no JSON result line"
        tail -40 "$STAGE/measure.out" >&2 || true
        exit 1
    fi

    # The measure line is already JSON, so the run object can be assembled
    # without any host-side interpreter.
    printf '{"hash":"%s","boot_seconds":%s,"kv_tokens":%s,"kv_concurrency":"%s","total_seconds":%s,"measure":%s}\n' \
        "$(tr -d '[:space:]' <"$STAGE/hash")" \
        "$boot_seconds" "${kv_tokens:-0}" "${kv_conc:-unknown}" \
        "$(( $(date +%s) - t_start ))" "$measure_json" >"$STAGE/run.json"

    log "bench window complete in $(( $(date +%s) - t_start ))s"
    # Memo only after a FULL success; the restore below still has to verify.
    exit 0
}

main "$@"
exit "$?"