#!/usr/bin/env bash
# SUM Parts - keep a phase script alive across anything that kills it
#
# The phase scripts already retry their own steps, and train_watchdog resumes
# training from its checkpoint. This is the layer above both: it relaunches the
# phase if the whole process disappears - a WSL VM restart, an OOM kill, a
# stray pkill.
#
# Safe because every phase is idempotent. A restart re-checks what is already
# done (conda env, patches, downloaded archives, training checkpoint) and
# continues from there rather than redoing it.
#
# Usage:
#   bash scripts/keepalive.sh setup      # phase A, no GPU needed
#   bash scripts/keepalive.sh train      # phase B, needs the GPU
#
# Unattended:
#   setsid nohup bash scripts/keepalive.sh setup > ~/keepalive.out 2>&1 &
#
# Stops when:
#   STATUS says DONE      -> success, exit 0
#   STATUS says FAILED    -> hard error (missing HF token, no GPU for phase B);
#                            retrying cannot fix it, exit 1
#   MAX_RESTARTS reached  -> exit 1
set -uo pipefail

SCRIPTS="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

TARGET="${1:-}"
case "$TARGET" in
    setup|run_setup) TARGET=setup; SCRIPT="$SCRIPTS/run_setup.sh"; OUT="$HOME/sum-parts/runs/setup" ;;
    train|run_train) TARGET=train; SCRIPT="$SCRIPTS/run_train.sh"; OUT="$HOME/sum-parts/runs/train" ;;
    *)
        echo "usage: bash keepalive.sh {setup|train}"
        echo
        echo "  setup  phase A - bootstrap, conda, CUDA extensions, data. No GPU needed."
        echo "  train  phase B - voxel_max measurement, training, evaluation. Needs the GPU."
        exit 2
        ;;
esac

STATUS="$OUT/STATUS"
KLOG="$OUT/keepalive.log"

MAX_RESTARTS="${MAX_RESTARTS:-40}"
COOLDOWN="${COOLDOWN:-90}"

mkdir -p "$OUT"

klog() { echo "[$(date '+%F %T')] $*" | tee -a "$KLOG"; }

state_of() {
    [ -f "$STATUS" ] || { echo "none"; return; }
    grep -E '^state' "$STATUS" | head -1 | cut -d: -f2- | tr -d ' '
}

klog "keepalive starting for phase '$TARGET' (max $MAX_RESTARTS restarts, ${COOLDOWN}s cooldown)"

restarts=0
while :; do
    s=$(state_of)
    case "$s" in
        DONE)
            klog "phase '$TARGET' reports DONE"
            [ "$TARGET" = setup ] && {
                klog "next, once the GPU is free:"
                klog "  setsid nohup bash $SCRIPTS/keepalive.sh train > ~/keepalive-train.out 2>&1 &"
            }
            exit 0
            ;;
        FAILED)
            klog "phase '$TARGET' reports FAILED -- restarting will not fix this:"
            sed 's/^/    /' "$STATUS" | tee -a "$KLOG"
            exit 1
            ;;
    esac

    # already running (started by hand, or by a previous loop)?
    if pgrep -f "$(basename "$SCRIPT")" | grep -qv "^$$\$"; then
        sleep 30
        continue
    fi

    if [ "$restarts" -ge "$MAX_RESTARTS" ]; then
        klog "hit MAX_RESTARTS=$MAX_RESTARTS -- stopping"
        exit 1
    fi

    if [ "$restarts" -gt 0 ]; then
        klog "not running (state='$s') -- restart #$restarts"
    else
        klog "launching $(basename "$SCRIPT")"
    fi

    bash "$SCRIPT"
    rc=$?
    klog "$(basename "$SCRIPT") exited rc=$rc, state='$(state_of)'"

    restarts=$((restarts + 1))
    sleep "$COOLDOWN"
done