#!/usr/bin/env bash
# 基于进度的卡死检测，取代之前"固定600秒墙钟时间"的误杀判定（2026-08-23）。
# 判定标准：denoise 的 step 数字长时间不变才算卡死，不看总耗时——大分辨率/长帧数任务本来
# 就该跑更久，不该被时间线误杀。
#
# 【2026-08-25修复】漏了一个阶段：采样跑完(step达到total_steps)之后还有VAE视频/音频解码+
# 存盘，这段没有独立进度信号，step会一直停在最终值不变——之前这段直接被并入
# STALL_LIMIT_RUNNING(10分钟)判定，实测至少一个1344×768长帧数案例这个阶段本身就超过608秒，
# 被误判"卡死"重启，把一批正常在跑/排队的任务(showcase案例批次)集体打断。单独给这个阶段
# 一个更宽的阈值，跟INIT阶段同样的思路——不同阶段没有可比的"正常耗时"，不能共用一个数字。
set -uo pipefail
BASE="http://127.0.0.1:8100"
STATE_FILE="/data/models/raylight/watchdog.state"
LOG_FILE="/data/models/raylight/watchdog.log"
STALL_LIMIT_RUNNING=600     # denoise 采样中(0<step<total_steps)，step 10分钟不变才算卡死
STALL_LIMIT_INIT=1200       # 还在FSDP初始化(step仍是0)阶段，给20分钟宽限
STALL_LIMIT_POSTPROC=1800   # 采样已跑完(step>=total_steps)，VAE解码/存盘阶段，给30分钟宽限
POLL_INTERVAL=20

log() { echo "$(date '+%F %T') $*" >> "$LOG_FILE"; }

# 状态文件格式：job_id step last_change_epoch
read_state() { cat "$STATE_FILE" 2>/dev/null || echo "none 0 0"; }
write_state() { echo "$1 $2 $3" > "$STATE_FILE"; }

log "watchdog started"

while true; do
  sleep "$POLL_INTERVAL"

  HEALTH=$(curl -s --max-time 10 "$BASE/health" 2>/dev/null)
  RUNNING=$(echo "$HEALTH" | python3 -c "import sys,json; print(json.load(sys.stdin).get('running',0))" 2>/dev/null)
  if [ "$RUNNING" != "1" ]; then
    write_state none 0 0
    continue
  fi

  JOBS=$(curl -s --max-time 10 "$BASE/v1/videos?limit=5" 2>/dev/null)
  CUR=$(echo "$JOBS" | python3 -c "
import sys, json
try:
    jobs = json.load(sys.stdin)
    for j in jobs:
        if j.get('status') == 'running':
            print(j['id']); break
except Exception:
    pass
" 2>/dev/null)
  [ -z "$CUR" ] && continue

  DETAIL=$(curl -s --max-time 10 "$BASE/v1/videos/$CUR" 2>/dev/null)
  STEP=$(echo "$DETAIL" | python3 -c "import sys,json; print(json.load(sys.stdin).get('step',0))" 2>/dev/null)
  [ -z "$STEP" ] && STEP=0
  TOTAL_STEPS=$(echo "$DETAIL" | python3 -c "import sys,json; print(json.load(sys.stdin).get('total_steps',0))" 2>/dev/null)
  [ -z "$TOTAL_STEPS" ] && TOTAL_STEPS=0

  read -r PREV_ID PREV_STEP PREV_TS < <(read_state)
  NOW=$(date +%s)

  if [ "$CUR" != "$PREV_ID" ] || [ "$STEP" != "$PREV_STEP" ]; then
    write_state "$CUR" "$STEP" "$NOW"
    log "job=$CUR step=$STEP (progress, reset stall timer)"
    continue
  fi

  STALLED_FOR=$((NOW - PREV_TS))
  if [ "$STEP" = "0" ]; then
    LIMIT=$STALL_LIMIT_INIT
  elif [ "$TOTAL_STEPS" != "0" ] && [ "$STEP" -ge "$TOTAL_STEPS" ] 2>/dev/null; then
    LIMIT=$STALL_LIMIT_POSTPROC
  else
    LIMIT=$STALL_LIMIT_RUNNING
  fi

  if [ "$STALLED_FOR" -ge "$LIMIT" ]; then
    log "job=$CUR step=$STEP 卡死判定触发：$STALLED_FOR 秒无进展（阈值 $LIMIT），执行重启"
    bash /data/models/raylight/start.sh >> "$LOG_FILE" 2>&1
    write_state none 0 0
    log "重启完成"
  fi
done