#!/bin/bash
#SBATCH --job-name=xvla-deploy
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=365-24:00:00
#SBATCH --partition=h100-private
#SBATCH --qos=high
#SBATCH --output=/tmp/slurm-%x-%j.out
#SBATCH --error=/tmp/slurm-%x-%j.err

set -euo pipefail

REPO_ROOT="/shared_work/xuan/project/x-vla-backup"
MODEL_PATH="${MODEL_PATH:-$REPO_ROOT/runnings/roboreal_all_20260428_g8_bs32.ckpt-100000}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8010}"
OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/runnings/deploy_${SLURM_JOB_ID}}"
DEVICE="${DEVICE:-cuda}"
PROCESSOR_PATH="${PROCESSOR_PATH:-}"
LORA_PATH="${LORA_PATH:-}"

resolve_model_path() {
    local raw_path="$1"
    local candidate=""

    if [[ -e "$raw_path" ]]; then
        printf '%s\n' "$raw_path"
        return 0
    fi

    # Accept the common shorthand "...run_name.ckpt-100000" by rewriting it to
    # the actual checkpoint directory layout "...run_name/ckpt-100000".
    if [[ "$raw_path" =~ ^(.+)\.ckpt-([0-9]+)$ ]]; then
        candidate="${BASH_REMATCH[1]}/ckpt-${BASH_REMATCH[2]}"
        if [[ -e "$candidate" ]]; then
            printf '%s\n' "$candidate"
            return 0
        fi
    fi

    return 1
}

choose_output_dir() {
    local requested_dir="$1"
    local candidate="$requested_dir"
    local suffix=""

    while [[ -f "$candidate/info.json" ]]; do
        suffix="$(date -u +%Y%m%dT%H%M%SZ)"
        if [[ -n "${SLURM_JOB_ID:-}" ]]; then
            suffix="${suffix}_job${SLURM_JOB_ID}"
        fi
        candidate="${requested_dir}_${suffix}"
        sleep 1
    done

    printf '%s\n' "$candidate"
}

cd "$REPO_ROOT"

export SHARED_ROOT="/shared_work/xuan"
export CONDA_BASE_PREFIX="$SHARED_ROOT/miniforge3"
export ENV_PREFIX="$SHARED_ROOT/conda-envs/xvla-stable"

source "$REPO_ROOT/scripts/activate_shared_env.sh"

if ! RESOLVED_MODEL_PATH="$(resolve_model_path "$MODEL_PATH")"; then
    echo "Model path not found: $MODEL_PATH" >&2
    if [[ "$MODEL_PATH" =~ \.ckpt-[0-9]+$ ]]; then
        echo "Hint: this script expects a checkpoint directory such as .../run_name/ckpt-100000." >&2
    fi
    exit 1
fi

MODEL_PATH="$RESOLVED_MODEL_PATH"

OUTPUT_DIR="$(choose_output_dir "$OUTPUT_DIR")"
mkdir -p "$OUTPUT_DIR"

export PYTHONUNBUFFERED=1

echo "Launching X-VLA deploy server"
echo "  MODEL_PATH=$MODEL_PATH"
echo "  HOST=$HOST"
echo "  PORT=$PORT"
echo "  OUTPUT_DIR=$OUTPUT_DIR"
echo "  DEVICE=$DEVICE"
echo "  CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}"
echo "  SLURM_JOB_ID=${SLURM_JOB_ID:-unset}"
echo "  SLURM_JOB_NODELIST=${SLURM_JOB_NODELIST:-unset}"

CMD=(
    python -m deploy
    --model_path "$MODEL_PATH"
    --host "$HOST"
    --port "$PORT"
    --output_dir "$OUTPUT_DIR"
    --device "$DEVICE"
    --disable_slurm
)

if [[ -n "$PROCESSOR_PATH" ]]; then
    CMD+=(--processor_path "$PROCESSOR_PATH")
fi

if [[ -n "$LORA_PATH" ]]; then
    CMD+=(--LoRA_path "$LORA_PATH")
fi

exec "${CMD[@]}"