#!/usr/bin/env bash
set -euo pipefail

if [[ $# -lt 2 || $# -gt 3 ]]; then
  echo "usage: $0 <uniform|residual> <run-id> [proposal.json]" >&2
  exit 64
fi

arm=$1
run_id=$2
proposal=${3:-}
project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
python="${OPTR_PYTHON:-python}"
config="${project}/configs/florence_bridge_20k_plus500.yaml"
accelerate_config="${project}/configs/deepspeed_florence_2gpu.yaml"

required_env=(OPTR_RUN_ROOT FLORENCE_MODEL_PATH BRIDGE_DATA_ROOT BRIDGE_STATISTICS_PATH BRIDGE_PROVENANCE_PATH BRIDGE_FINGERPRINT_PATH BRIDGE_MANIFEST_PATH COMMON_START_20K)
for name in "${required_env[@]}"; do
  [[ -n "${!name:-}" ]] || { echo "missing environment variable: $name" >&2; exit 64; }
done
[[ -d "$COMMON_START_20K" ]] || { echo "common start not found: $COMMON_START_20K" >&2; exit 66; }
[[ -f "$COMMON_START_20K/checkpoint.json" ]] || { echo "checkpoint.json missing under common start" >&2; exit 66; }
[[ -f "$BRIDGE_STATISTICS_PATH" ]] || { echo "statistics file missing: $BRIDGE_STATISTICS_PATH" >&2; exit 66; }
[[ -f "$BRIDGE_PROVENANCE_PATH" ]] || { echo "provenance file missing: $BRIDGE_PROVENANCE_PATH" >&2; exit 66; }
[[ -f "$BRIDGE_FINGERPRINT_PATH" ]] || { echo "content fingerprint missing: $BRIDGE_FINGERPRINT_PATH" >&2; exit 66; }
[[ -f "$BRIDGE_MANIFEST_PATH" ]] || { echo "anchor manifest missing: $BRIDGE_MANIFEST_PATH" >&2; exit 66; }

case "$arm" in
  uniform)
    [[ -z "$proposal" ]] || { echo "uniform does not accept a proposal" >&2; exit 64; }
    entry=starVLA/training/train_starvla.py
    proposal_args=(--datasets.vla_data.temporal_proposal_path=null)
    ;;
  residual)
    [[ -f "$proposal" ]] || { echo "proposal not found: $proposal" >&2; exit 66; }
    entry=starVLA/training/train_starvla.py
    proposal_args=(--datasets.vla_data.temporal_proposal_path="$proposal")
    ;;
  *) echo "unknown arm: $arm" >&2; exit 64 ;;
esac

cd "$project"
export PYTHONPATH="${project}:${PYTHONPATH:-}"
export WANDB_MODE="${WANDB_MODE:-offline}"
export PYTHONUNBUFFERED=1
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
mkdir -p "$OPTR_RUN_ROOT"
"$python" -m accelerate.commands.launch \
  --config_file "$accelerate_config" \
  --num_processes 2 \
  "$entry" \
  --config_yaml "$config" \
  --run_id "$run_id" \
  --run_root_dir "$OPTR_RUN_ROOT" \
  "${proposal_args[@]}"