#!/usr/bin/env bash
# SUM Parts - keep run_all.sh alive across anything that kills it
#
# run_all.sh already retries individual phases, and train_watchdog resumes
# training from its checkpoint. This is the layer above both: it restarts
# run_all itself if the whole process disappears -- a WSL VM restart, an OOM
# kill, a stray pkill.
#
# That is 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.
#
# Stops when:
#   - STATUS says DONE           -> success, exits 0
#   - STATUS says FAILED         -> a hard error like a missing HF token;
#                                   retrying cannot fix it, exits 1
#   - MAX_RESTARTS reached       -> exits 1
#
# Usage (this is the one command to run before leaving):
#   setsid nohup bash scripts/keepalive.sh > ~/keepalive.out 2>&1 &
#
# Check on it:
#   cat ~/sum-parts/runs/run_all/STATUS
#   tail -f ~/sum-parts/runs/run_all/run.log
set -uo pipefail

SCRIPTS="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUT="$HOME/sum-parts/runs/run_all"
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 (max $MAX_RESTARTS restarts, ${COOLDOWN}s cooldown)"

restarts=0
while :; do
    s=$(state_of)
    case "$s" in
        DONE)
            klog "run_all reports DONE -- finished"
            exit 0
            ;;
        FAILED)
            klog "run_all reports FAILED -- a hard error that restarting will not fix:"
            sed 's/^/    /' "$STATUS" | tee -a "$KLOG"
            exit 1
            ;;
    esac

    if pgrep -f "run_all.sh" | 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 "run_all is not running (state='$s') -- restart #$restarts"
    else
        klog "launching run_all"
    fi

    bash "$SCRIPTS/run_all.sh"
    rc=$?
    klog "run_all exited rc=$rc, state='$(state_of)'"

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