#!/bin/bash
set -euo pipefail

# Submit and optionally wait for a short smoke-test run.
# Usage:
#   slurm/run_gguf_smoke.sh
#   WAIT_FOR_COMPLETION=0 slurm/run_gguf_smoke.sh

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:20:00}"
SMOKE_TEST_MAX_CASES="${SMOKE_TEST_MAX_CASES:-3}"
SMOKE_TEST_EVAL_K="${SMOKE_TEST_EVAL_K:-1}"
SMOKE_TEST_CONTEXT_SIZE="${SMOKE_TEST_CONTEXT_SIZE:-2048}"
WAIT_FOR_COMPLETION="${WAIT_FOR_COMPLETION:-1}"
SMOKE_RUN_ID="${SMOKE_RUN_ID:-smoke-$(date +%Y%m%dT%H%M%S)}"
SMOKE_XDG_RUNTIME_DIR="${SMOKE_XDG_RUNTIME_DIR:-/tmp/xdg-runtime-${UID:-0}}"
SMOKE_SERVER_BACKEND="${SMOKE_SERVER_BACKEND:-${SERVER_BACKEND:-vllm}}"

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

cd "$REPO_DIR"

if [[ -x "slurm/check_env_gguf_eval.sh" ]]; then
    echo "Running preflight..."
    slurm/check_env_gguf_eval.sh
fi

echo "Submitting smoke test job..."
echo "Smoke backend: ${SMOKE_SERVER_BACKEND}"
JOB_ID="$(sbatch --parsable \
    --time="$SMOKE_TIME" \
    --export=ALL,SERVER_BACKEND="$SMOKE_SERVER_BACKEND",SMOKE_TEST_MODE=1,EVAL_RUN_ID="$SMOKE_RUN_ID",EVAL_MAX_CASES="$SMOKE_TEST_MAX_CASES",EVAL_K="$SMOKE_TEST_EVAL_K",CONTEXT_SIZE="$SMOKE_TEST_CONTEXT_SIZE",XDG_RUNTIME_DIR="$SMOKE_XDG_RUNTIME_DIR",BASH_ENV= \
    slurm/eval_gguf_eval.sh)"

echo "Smoke job submitted: ${JOB_ID}"
echo "Smoke run id: ${SMOKE_RUN_ID}"

echo "Logs:"
echo "  eval-${JOB_ID}.out"

echo ""
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
else
    echo "WARN: squeue not available; cannot block on live queue state."
    sleep 5
fi

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)"
else
    STATE="UNKNOWN"
    EXIT_CODE="UNKNOWN"
fi

echo "Smoke job state: ${STATE} (ExitCode=${EXIT_CODE})"

if [[ "${STATE^^}" != COMPLETED* ]]; then
    echo "Smoke test failed or did not complete cleanly."
    echo "Check: tail -n 200 eval-${JOB_ID}.out"
    exit 1
fi

echo "Smoke test passed."
