#!/usr/bin/env bash
# run_h3_supervisor.sh — H3 批次「自愈续跑」：机器重启/ComfyUI 掉线都能自动拉起来继续跑，直到全部出齐。
#
# 背景：GPU 机在 H3 长时满载下会随机整机硬挂；`kernel.panic=10` 会触发自动重启（09-25 实测有效），
# 但 ComfyUI 不会自启 → 本脚本负责「等机器 → 拉起 ComfyUI/RAM 哨兵 → 断点续跑」。
#
# 用法：setsid nohup bash run_h3_supervisor.sh >> logs/h3_supervisor.log 2>&1 < /dev/null &
set -uo pipefail
cd "$(dirname "$0")" || exit 1
GPU=192.168.31.31
SSH=~/Downloads/ssh_exec.sh
COOLDOWN="${COOLDOWN:-300}"
SIZE="${SIZE:-1344x768}"
MAX_ROUNDS="${MAX_ROUNDS:-240}"
COMFY_CMD='cd /home/zyw/ComfyUI && setsid nohup env PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 venv/bin/python main.py --listen 0.0.0.0 --port 8189 --lowvram --force-fp16 --reserve-vram 1.0 --fast-disk --cache-lru 2 --disable-pinned-memory --preview-method none --enable-manager > /home/zyw/comfyui_8189_opt.log 2>&1 < /dev/null & sleep 2; setsid nohup bash /home/zyw/ram_guard.sh > /dev/null 2>&1 < /dev/null &'

log() { echo "[$(date '+%m-%d %H:%M:%S')] $*"; }
alive() { ping -c 1 -W 2 "$GPU" >/dev/null 2>&1; }
comfy_up() { timeout 30 bash "$SSH" 'curl -s --noproxy "*" -m 8 -o /dev/null -w "%{http_code}" http://127.0.0.1:8189/system_stats' 2>/dev/null | grep -q 200; }

todo_count() {
  python3 - <<'PY'
import json, os
items = json.load(open("prompts/h3/batch.json", encoding="utf-8"))
print(sum(1 for i in items if not os.path.exists(f"06-H3片段/{i['shot']}/{i['shot']}.mp4")))
PY
}
total_count() { python3 -c "import json;print(len(json.load(open('prompts/h3/batch.json',encoding='utf-8'))))"; }

for round in $(seq 1 "$MAX_ROUNDS"); do
  todo=$(todo_count); total=$(total_count)
  log "第 $round 轮：已完成 $(( total - todo ))/$total，待跑 $todo"
  [ "$todo" -eq 0 ] && { log "全部完成 ✅"; exit 0; }

  if ! alive; then
    log "GPU 机不可达（硬挂中，等 kernel.panic 自动重启或人工断电）——60s 后重试"
    sleep 60; continue
  fi

  if ! comfy_up; then
    log "机器在线但 ComfyUI 未就绪 → 自动拉起 ComfyUI + RAM 哨兵"
    timeout 90 bash "$SSH" "$COMFY_CMD" >/dev/null 2>&1
    for i in $(seq 1 20); do
      sleep 20
      if comfy_up; then log "ComfyUI 就绪（等待 $((i*20))s）"; break; fi
    done
    if ! comfy_up; then log "ComfyUI 仍未起来，下一轮再试"; continue; fi
  fi

  log "开始续跑（size=$SIZE cooldown=$COOLDOWN）"
  python3 -u run_h3_batch.py --size "$SIZE" --cooldown "$COOLDOWN" 2>&1 | tail -40

  if ! alive; then
    log "批次中断且机器不可达；等 5 分钟再探"
    sleep 300
  else
    sleep 60
  fi
done
log "达到最大轮数，退出"