#!/bin/bash
# Watchdog: detect failed jobs in last 30min and resubmit if recoverable
# Run every 30min; idempotent (guards against double-submit via marker files)

set -e
LOGDIR=/arf/scratch/stakan/hitit-proje/logs/h100
MARKERS=/arf/scratch/stakan/hitit-proje/logs/watchdog_markers
mkdir -p $MARKERS

SLURM=/arf/scratch/stakan/hitit-proje/hitit_ocr/scripts/pipeline_h100
NOW=$(date +%s)

# Get failures in last 45min
SINCE=$(date -d '45 minutes ago' +%Y-%m-%dT%H:%M:%S)
sacct -u stakan --starttime=$SINCE --format=JobID,JobName%25,State 2>&1 \
    | grep -v "\." \
    | awk '$3=="FAILED"{print $1" "$2}' > $MARKERS/recent_fails.txt

if [ ! -s $MARKERS/recent_fails.txt ]; then
    echo "$(date) No recent fails"; exit 0
fi

echo "$(date) Recent fails:"
cat $MARKERS/recent_fails.txt

# Current queue (pending+running) — skip resubmit if job already present
CURRENT_NAMES=$(squeue -u stakan -h -o "%j" 2>/dev/null | sort -u)

# Deduplicate fails by job name (keep most recent)
awk '{names[$2]=$1} END{for(n in names)print names[n]" "n}' $MARKERS/recent_fails.txt > $MARKERS/fails_dedup.txt
mv $MARKERS/fails_dedup.txt $MARKERS/recent_fails.txt

# For each, attempt resubmit (guarded)
while read JID JNAME; do
    # Skip if same job already in queue (prevents duplicate submission)
    if echo "$CURRENT_NAMES" | grep -Fxq "$JNAME"; then
        echo "  $JNAME: already in queue, skip"; continue
    fi
    # Guard: don't resubmit same job >3 times
    COUNT_FILE=$MARKERS/resubmit_count_${JNAME}.txt
    COUNT=$(cat $COUNT_FILE 2>/dev/null || echo 0)
    if [ "$COUNT" -ge 3 ]; then
        echo "  $JNAME: max retries reached, skipping"
        continue
    fi

    # Find matching SLURM script by jobname prefix
    SCRIPT=""
    case "$JNAME" in
        *"ssl_v3L"*)       SCRIPT=$SLURM/60_ssl_dinov3l.slurm ;;
        *"cRT"*)           SCRIPT=$SLURM/51_crt_decoupled.slurm ;;
        *"distill_v4"*)    SCRIPT=$SLURM/44_distill_v4.slurm ;;
        *"pseudo_cls"*)    SCRIPT=$SLURM/25_pseudo_cls.slurm ;;
        *"active"*)        SCRIPT=$SLURM/65_active_export.slurm ;;
        *"yolo_fold"*)     SCRIPT=$SLURM/05_yolo_5fold.slurm ;;
        *"ensemble_v4"*)   SCRIPT=$SLURM/91_ensemble_v4_fixed.slurm ;;
        *"mega"*)          SCRIPT=$SLURM/92_mega_ensemble.slurm ;;
        *"postproc"*|*"100_pp"*) SCRIPT=$SLURM/100_postproc_all.slurm ;;
        *"101_adv"*)       SCRIPT=$SLURM/101_advanced_batch.slurm ;;
        *"swin_v4"*|*"34_swin"*)       SCRIPT=$SLURM/34_swin_v4.slurm ;;
        *"swin_v6"*|*"63_swin"*)       SCRIPT=$SLURM/63_swin_v6.slurm ;;
        *"dv3_v7"*|*"80_dv3"*)   SCRIPT=$SLURM/80_dinov3_v7_ssl_ft.slurm ;;
        *"v7anchor"*)      SCRIPT=$SLURM/81_dinov3_v7_anchor.slurm ;;
        *"dv3_336"*)       SCRIPT=$SLURM/64_dinov3_336_v6.slurm ;;
        *"focalSAM"*)      SCRIPT=$SLURM/53_focal_sam_dinov3.slurm ;;
        *"90_methods"*)    SCRIPT=$SLURM/90_methods_batch.slurm ;;
        *"110_ultra"*)     SCRIPT=$SLURM/110_ultra_ensemble.slurm ;;
        *"120_ps2"*)       SCRIPT=$SLURM/120_pseudo_iter2.slurm ;;
        *"130_v8"*)        SCRIPT=$SLURM/130_dinov3_v8.slurm ;;
        *"131_rand"*)      SCRIPT=$SLURM/131_dinov3_randsplit.slurm ;;
        *"140_ensv8"*)     SCRIPT=$SLURM/140_ensemble_v8.slurm ;;
        *"150_ps2m"*)      SCRIPT=$SLURM/150_pseudo_iter2_multi.slurm ;;
        *"151_v9"*)        SCRIPT=$SLURM/151_dinov3_v9.slurm ;;
        *"160_seed"*)      SCRIPT=$SLURM/160_seed_ensemble.slurm ;;
        *"170_xval_rs"*)   SCRIPT=$SLURM/170_crossval_randsplit.slurm ;;
        *"171_rs_eval"*)   SCRIPT=$SLURM/171_randsplit_only_eval.slurm ;;
        *"180_v9pp"*)      SCRIPT=$SLURM/180_v9_postproc.slurm ;;
        *"190_seedeval"*)  SCRIPT=$SLURM/190_seed_ensemble_eval.slurm ;;
        *"200_super"*)     SCRIPT=$SLURM/200_super_mega.slurm ;;
        *"220_tlh_ut"*)    SCRIPT=$SLURM/220_tlhdig_utils_build.slurm ;;
        *"221_signLM"*)    SCRIPT=$SLURM/221_sign_lm.slurm ;;
        *"223_weak"*)      SCRIPT=$SLURM/223_weak_sup.slurm ;;
        *"224_posth"*)     SCRIPT=$SLURM/224_post_heads_batch.slurm ;;
        *"230_v10cs"*)     SCRIPT=$SLURM/230_dinov3_v10_crosssrc.slurm ;;
        *"231_v11mf"*)     SCRIPT=$SLURM/231_v11_manifest.slurm ;;
        *"232_v11"*)       SCRIPT=$SLURM/232_dinov3_v11_full.slurm ;;
        *"240_v12mf"*)     SCRIPT=$SLURM/240_build_v12.slurm ;;
        *"242_v12"*)       SCRIPT=$SLURM/242_dinov3_v12_full.slurm ;;
        *"250_v13a"*)      SCRIPT=$SLURM/250_v13a_dinov3l.slurm ;;
        *"251_v13b"*)      SCRIPT=$SLURM/251_v13b_convnext.slurm ;;
        *"252_v13c"*)      SCRIPT=$SLURM/252_v13c_dinov3b.slurm ;;
        *"253_v13d"*)      SCRIPT=$SLURM/253_v13d_swin.slurm ;;
        *)
            echo "  $JNAME: no script mapping, skip"; continue ;;
    esac

    if [ ! -f "$SCRIPT" ]; then
        echo "  $JNAME: script $SCRIPT missing"; continue
    fi

    # Special handling: remove stale distill cache
    if [[ "$JNAME" == *"distill_v4"* ]]; then
        rm -f /arf/scratch/stakan/hitit-proje/hitit_ocr/runs/h100/distilled_student_v4/teacher_cache.pt
    fi

    NEW_JID=$(sbatch --parsable $SCRIPT 2>&1)
    if [[ "$NEW_JID" =~ ^[0-9]+$ ]]; then
        echo "$((COUNT+1))" > $COUNT_FILE
        echo "  $JNAME resubmitted as $NEW_JID (retry $((COUNT+1)))"
    else
        echo "  $JNAME resubmit failed: $NEW_JID"
    fi
done < $MARKERS/recent_fails.txt

# Summary
echo
echo "=== Current queue ==="
squeue -u stakan -o "%.10i %.22j %.10T %.8M" | head -15