#!/bin/bash
# =============================================================================
# GPU Watchdog for AMD BC-250 (Cyan Skillfish)
# =============================================================================
# Monitors kernel logs for KIQ fence timeouts and immediately kills the
# offending GPU process before the timeout cascade crashes the system.
#
# BC-250 Crash Pattern:
#   1st KIQ timeout → 12s → 2nd KIQ timeout → 12s → cascade → hard crash
#   Window to act: ~10 seconds after first timeout
#
# Usage:
#   ./gpu_watchdog.sh                    # Monitor and auto-kill
#   ./gpu_watchdog.sh --dry-run          # Monitor only, don't kill
#   ./gpu_watchdog.sh --max-kiq 2        # Kill after 2 KIQ timeouts
# =============================================================================

set -euo pipefail

MAX_KIQ_TIMEOUTS=1
CHECK_INTERVAL=2
DRY_RUN=false
LOG_FILE="/tmp/gpu_watchdog.log"

while [[ $# -gt 0 ]]; do
    case "$1" in
        --dry-run)    DRY_RUN=true; shift ;;
        --max-kiq)    MAX_KIQ_TIMEOUTS="$2"; shift 2 ;;
        --interval)   CHECK_INTERVAL="$2"; shift 2 ;;
        *)            echo "Unknown: $1"; exit 1 ;;
    esac
done

log() {
    local msg="[$(date '+%H:%M:%S')] $1"
    echo "$msg"
    echo "$msg" >> "$LOG_FILE"
}

get_kiq_count() {
    local count
    count=$(journalctl -k -b --no-pager 2>/dev/null | grep -c "timeout waiting for kiq fence" | head -1 | tr -d '[:space:]')
    echo "${count:-0}"
}

kill_gpu_processes() {
    local pids
    pids=$(fuser /dev/kfd 2>/dev/null || pgrep -f "sdcpp-restapi" 2>/dev/null || true)
    pids=$(echo "$pids" | xargs)
    
    if [[ -z "$pids" ]]; then
        log "WARN: No GPU processes found"
        return 1
    fi
    
    for pid in $pids; do
        local name
        name=$(ps -p "$pid" -o comm= 2>/dev/null || echo "?")
        if [[ "$DRY_RUN" == "true" ]]; then
            log "DRY-RUN: Would SIGKILL PID $pid ($name)"
        else
            log "SIGKILL PID $pid ($name)"
            kill -9 "$pid" 2>/dev/null || true
        fi
    done
    
    if [[ "$DRY_RUN" == "false" ]]; then
        systemctl --user stop zimage 2>/dev/null || true
        log "Stopped zimage service"
    fi
}

echo "=========================================="
echo " GPU Watchdog — AMD BC-250"
echo " Kill after: $MAX_KIQ_TIMEOUTS KIQ timeout(s)"
echo " Interval:   ${CHECK_INTERVAL}s"
echo " Dry run:    $DRY_RUN"
echo "=========================================="

log "Watchdog started"

INITIAL_KIQ=$(get_kiq_count)
PREV_KIQ=$INITIAL_KIQ
log "Baseline KIQ count: $INITIAL_KIQ"

[[ "$INITIAL_KIQ" -gt 0 ]] && log "WARNING: GPU already has $INITIAL_KIQ pre-existing KIQ timeouts"

while true; do
    CUR=$(get_kiq_count)
    NEW=$((CUR - INITIAL_KIQ))
    
    if [[ "$CUR" -ne "$PREV_KIQ" ]]; then
        log "!! KIQ timeout #$CUR (new=$NEW)"
        
        if [[ "$NEW" -ge "$MAX_KIQ_TIMEOUTS" ]]; then
            log "!!! THRESHOLD HIT — EMERGENCY KILL !!!"
            kill_gpu_processes
            log "GPU processes killed. Reboot needed for clean GPU."
            INITIAL_KIQ=$CUR
            log "Watchdog reset, continuing..."
        fi
    fi
    
    PREV_KIQ=$CUR
    sleep "$CHECK_INTERVAL"
done