#!/bin/bash
# ============================================================
# Dynamic SLURM job submitter
# Auto-detects account, tries multiple partitions
#
# Usage:
#   bash chengyue_script/dvlm/utils/submit.sh <script.sh> [extra sbatch args...]
#
# Examples:
#   bash chengyue_script/dvlm/utils/submit.sh my_train.sh
#   bash chengyue_script/dvlm/utils/submit.sh my_train.sh -N 4
#   bash chengyue_script/dvlm/utils/submit.sh my_train.sh --partition=batch_block3
#
# Environment variables:
#   SLURM_PARTITIONS  - comma-separated partitions (default: auto-detect all batch_block*)
#   SLURM_ACCOUNT     - override account (default: auto-detect)
# ============================================================

set -euo pipefail

if [ $# -lt 1 ]; then
    echo "Usage: bash $0 <script.sh> [extra sbatch args...]"
    exit 1
fi

SCRIPT="$1"
shift

if [ ! -f "$SCRIPT" ]; then
    echo "ERROR: Script not found: $SCRIPT"
    exit 1
fi

# Auto-detect SLURM account
if [ -z "${SLURM_ACCOUNT:-}" ]; then
    SLURM_ACCOUNT=$(sacctmgr -nP show assoc where user=$(whoami) format=account | head -1)
fi
if [ -z "$SLURM_ACCOUNT" ]; then
    echo "ERROR: Could not detect SLURM account for user $(whoami)"
    exit 1
fi

# Auto-detect available GPU partitions (batch_block* and backfill_block*)
# Filter out partitions with 0 total nodes
if [ -z "${SLURM_PARTITIONS:-}" ]; then
    SLURM_PARTITIONS=$(sinfo -h -o "%P %D" | sed 's/\*\b//' | awk '$2 > 0 {print $1}' | grep -E '^(batch_block|backfill_block)[0-9]+$' | paste -sd, -)
fi
if [ -z "$SLURM_PARTITIONS" ]; then
    echo "ERROR: No available partitions found"
    exit 1
fi

echo "Account:    ${SLURM_ACCOUNT}"
echo "Partitions: ${SLURM_PARTITIONS}"
echo "Script:     ${SCRIPT}"
echo "Extra args: $*"
echo "---"

sbatch --account="${SLURM_ACCOUNT}" --partition="${SLURM_PARTITIONS}" "$@" "${SCRIPT}"
