#!/usr/bin/env bash
# comfyui-stall-watchdog.sh —— 生成中途卡死检测（GPU 机 zywpc）
#
# 背景：2026-09-15 18:38 CST 硬挂时，机器整机失去响应 22 分 42 秒才被人为复位；
#       systemd 的硬件 watchdog 当时是关的（RuntimeWatchdogSec=off，/dev/watchdog 不存在）。
# 本脚本是**用户态**兜底：生成任务在跑但进度长时间不推进 / nvidia-smi 失联时，
#       记录证据，并可选择直接重启机器（默认只记录，不重启）。
#
# 用法（默认：只记录，不重启）
#   bash ~/comfyui-stall-watchdog.sh            # 前台
#   REBOOT_ON_STALL=1 bash ~/comfyui-stall-watchdog.sh   # 卡死自动 systemctl reboot
# 建议交给 systemd timer 或 cron 每分钟跑一次，而不是常驻。

set -u

LOG="${STALL_LOG:-/home/zyw/comfyui-stall.log}"
STATE="${STALL_STATE:-/run/comfyui-stall.state}"
COMFY_LOG="${COMFY_LOG:-/home/zyw/comfyui_8189_opt.log}"
STALL_SECONDS="${STALL_SECONDS:-600}"     # 进度行 10 分钟不动 = 判定卡死
REBOOT_ON_STALL="${REBOOT_ON_STALL:-0}"   # 1 = 卡死即重启
API="http://127.0.0.1:8189"

log() { printf '%s %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$*" >>"$LOG"; }

now=$(date +%s)

# 1) 有没有任务在跑
running=0
if command -v curl >/dev/null 2>&1; then
  prompt=$(timeout 8 curl -s --noproxy '*' "$API/prompt" 2>/dev/null)
  case "$prompt" in
    *'"queue_remaining": 0'*) running=0 ;;
    *'"queue_remaining"'*)    running=1 ;;
    *)                        running=0 ;;   # API 都没回应，交给下面的 GPU 检测
  esac
fi

# 2) GPU 是否还在
gpu_ok=0
if timeout 15 nvidia-smi -L >/dev/null 2>&1; then gpu_ok=1; fi

# 3) 进度行指纹（进度条每步都会刷新）
sig=""
if [ -f "$COMFY_LOG" ]; then
  sig=$(tail -c 400 "$COMFY_LOG" 2>/dev/null | tr '\r' '\n' | grep -aE "it/s|s/it" | tail -1)
fi

prev_sig=""; prev_ts=0
if [ -f "$STATE" ]; then
  # 格式：<epoch> <signature>
  prev_ts=$(cut -d' ' -f1 "$STATE" 2>/dev/null || echo 0)
  prev_sig=$(cut -d' ' -f2- "$STATE" 2>/dev/null || echo "")
fi

if [ "$gpu_ok" -eq 0 ]; then
  log "GPU-LOST nvidia-smi 无响应（running=$running）"
  if [ "$REBOOT_ON_STALL" = "1" ]; then log "REBOOT: GPU 失联，执行 systemctl reboot"; systemctl reboot; fi
  exit 0
fi

if [ "$running" -eq 1 ]; then
  if [ -n "$sig" ] && [ "$sig" = "$prev_sig" ] && [ "$prev_ts" -gt 0 ]; then
    stalled=$(( now - prev_ts ))
    if [ "$stalled" -ge "$STALL_SECONDS" ]; then
      log "STALL 生成中但进度行 $stalled 秒未变化: $sig"
      if [ "$REBOOT_ON_STALL" = "1" ]; then log "REBOOT: 判定卡死，执行 systemctl reboot"; systemctl reboot; fi
    fi
  fi
fi

printf '%s %s\n' "$now" "$sig" >"$STATE" 2>/dev/null || true