#!/bin/bash
# ============================================================
# Auto-submit library - source this at the top of any SLURM script
#
# When not inside a SLURM job, auto-detects account & partitions
# and resubmits the calling script via sbatch, then exits.
# When inside a SLURM job, does nothing (script continues normally).
#
# Usage: Add these 2 lines near the top of your script (after #SBATCH lines):
#
#   LMFLOW_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/relative/path/to/LMFlow" && pwd)"
#   source "${LMFLOW_ROOT}/chengyue_script/dvlm/utils/auto_submit.sh"
#
# Environment variables (optional overrides):
#   SLURM_ACCOUNT     - override account (default: auto-detect)
#   SLURM_PARTITIONS  - comma-separated partitions (default: auto-detect)
# ============================================================

# Only run auto-submit logic if NOT already inside a SLURM job
if [ -z "${SLURM_JOB_ID:-}" ]; then
    # 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)" >&2
        exit 1
    fi

    # Auto-detect available GPU partitions, filter out empty ones
    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" >&2
        exit 1
    fi

    # The calling script is BASH_SOURCE[1] (this file is BASH_SOURCE[0])
    CALLER_SCRIPT="${BASH_SOURCE[1]}"

    echo "=== Auto-submit ==="
    echo "Account:    ${SLURM_ACCOUNT}"
    echo "Partitions: ${SLURM_PARTITIONS}"
    echo "Script:     ${CALLER_SCRIPT}"
    echo "==================="

    sbatch --account="${SLURM_ACCOUNT}" --partition="${SLURM_PARTITIONS}" "${CALLER_SCRIPT}" "$@"
    exit $?
fi
