#!/bin/bash
set -euo pipefail

# Submit and optionally wait for SWE-Bench Lite smoke generation job.

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd -P)"
REPO_DIR="$(cd "$SCRIPT_DIR/.." && pwd -P)"
if [[ -f "${SCRIPT_DIR}/_load_local_env.sh" ]]; then
    # shellcheck source=slurm/_load_local_env.sh
    source "${SCRIPT_DIR}/_load_local_env.sh"
    load_slurm_local_env "$SCRIPT_DIR"
fi

SMOKE_TIME="${SMOKE_TIME:-00:45:00}"
SBL_SMOKE_MAX_INSTANCES="${SBL_SMOKE_MAX_INSTANCES:-10}"
SBL_SMOKE_CONTEXT_SIZE="${SBL_SMOKE_CONTEXT_SIZE:-2048}"
SBL_EVAL_JOBS="${SBL_EVAL_JOBS:-1}"
WAIT_FOR_COMPLETION="${WAIT_FOR_COMPLETION:-1}"
PREPARE_HANDOFF_AFTER_SMOKE="${PREPARE_HANDOFF_AFTER_SMOKE:-1}"
SBL_RUN_ID="${SBL_RUN_ID:-swebench-smoke-$(date +%Y%m%dT%H%M%S)}"

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

cd "$REPO_DIR"

if [[ -x "slurm/check_env_swebench_lite.sh" ]]; then
    echo "Running SWE-Bench Lite preflight..."
    slurm/check_env_swebench_lite.sh
fi

echo "Submitting SWE-Bench Lite smoke job..."
JOB_ID="$(sbatch --parsable \
    --time="$SMOKE_TIME" \
    --export=ALL,SBL_SMOKE_MODE=1,SBL_RUN_ID="$SBL_RUN_ID",SBL_SMOKE_MAX_INSTANCES="$SBL_SMOKE_MAX_INSTANCES",SBL_SMOKE_CONTEXT_SIZE="$SBL_SMOKE_CONTEXT_SIZE",SBL_EVAL_JOBS="$SBL_EVAL_JOBS" \
    slurm/eval_swebench_lite_generate.sh)"

echo "Smoke job submitted: ${JOB_ID}"
echo "Smoke run id: ${SBL_RUN_ID}"
echo "Log: swebench-lite-gen-${JOB_ID}.out"

if (( WAIT_FOR_COMPLETION == 0 )); then
    echo "Not waiting for completion (WAIT_FOR_COMPLETION=0)."
    exit 0
fi

echo "Waiting for smoke job ${JOB_ID} to finish..."
if command -v squeue >/dev/null 2>&1; then
    while squeue -h -j "$JOB_ID" 2>/dev/null | grep -q .; do
        sleep 10
    done
fi

STATE="UNKNOWN"
EXIT_CODE="UNKNOWN"
if command -v sacct >/dev/null 2>&1; then
    STATE="$(sacct -n -X -j "$JOB_ID" --format=State | head -n 1 | xargs || true)"
    EXIT_CODE="$(sacct -n -X -j "$JOB_ID" --format=ExitCode | head -n 1 | xargs || true)"
fi

echo "Smoke job state: ${STATE} (ExitCode=${EXIT_CODE})"
if [[ "${STATE^^}" != COMPLETED* ]]; then
    echo "Smoke run failed or did not complete cleanly."
    echo "Check: tail -n 200 swebench-lite-gen-${JOB_ID}.out"
    exit 1
fi

if (( PREPARE_HANDOFF_AFTER_SMOKE == 1 )); then
    echo "Running post-smoke handoff validation..."
    python -m evals.swebench.cli prepare-handoff --run-id "$SBL_RUN_ID"
    echo "Post-smoke handoff validation passed."
else
    echo "Skipping post-smoke handoff validation (PREPARE_HANDOFF_AFTER_SMOKE=0)."
fi

echo "Smoke test passed."
