#!/bin/bash
set -euo pipefail

# Submit dependent SWE-Bench Lite generation jobs sharing a single run_id.
# Usage:
#   slurm/submit_swebench_lite_resume_chain.sh <RUN_ID> [CHAIN_LENGTH]

if [[ $# -lt 1 || $# -gt 2 ]]; then
    echo "Usage: $0 <RUN_ID> [CHAIN_LENGTH]"
    exit 2
fi

RUN_ID="$1"
CHAIN_LENGTH="${2:-3}"
CHAIN_TIME="${CHAIN_TIME:-12:00:00}"
DEPENDENCY_TYPE="${DEPENDENCY_TYPE:-afterany}"
REQUIRE_SMOKE="${REQUIRE_SMOKE:-0}"
SBL_EVAL_JOBS="${SBL_EVAL_JOBS:-1}"

if [[ "${REQUIRE_SMOKE,,}" =~ ^(true|yes|on)$ ]]; then
    REQUIRE_SMOKE=1
elif [[ "${REQUIRE_SMOKE,,}" =~ ^(false|no|off)$ ]]; then
    REQUIRE_SMOKE=0
fi

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
REPO_DIR="$(cd "$SCRIPT_DIR/.." && pwd -P)"

cd "$REPO_DIR"

if (( REQUIRE_SMOKE == 1 )); then
    echo "Running smoke test before resume-chain submissions..."
    WAIT_FOR_COMPLETION=1 slurm/run_swebench_lite_smoke.sh
fi

echo "Submitting SWE-Bench Lite resume chain"
echo "run_id=${RUN_ID}"
echo "chain_length=${CHAIN_LENGTH}"
echo "time=${CHAIN_TIME}"
echo "dependency_type=${DEPENDENCY_TYPE}"
printf '%-8s %-14s %-14s\n' "INDEX" "JOB_ID" "DEPENDS_ON"
printf '%-8s %-14s %-14s\n' "-----" "------" "----------"

prev_job=""
for ((i=1; i<=CHAIN_LENGTH; i++)); do
    if [[ -z "$prev_job" ]]; then
        job_id="$(sbatch --parsable \
            --time="$CHAIN_TIME" \
            --export=ALL,SBL_RUN_ID="$RUN_ID",SBL_RESUME=1,SBL_EVAL_JOBS="$SBL_EVAL_JOBS" \
            slurm/eval_swebench_lite_generate.sh)"
        dep="-"
    else
        job_id="$(sbatch --parsable \
            --dependency="${DEPENDENCY_TYPE}:${prev_job}" \
            --time="$CHAIN_TIME" \
            --export=ALL,SBL_RUN_ID="$RUN_ID",SBL_RESUME=1,SBL_EVAL_JOBS="$SBL_EVAL_JOBS" \
            slurm/eval_swebench_lite_generate.sh)"
        dep="$prev_job"
    fi

    printf '%-8s %-14s %-14s\n' "$i" "$job_id" "$dep"
    prev_job="$job_id"
done
