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

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
WORKSPACE_ROOT="${WORKSPACE_ROOT:-/workspace}"
RLBENCH_ENV="${WORKSPACE_ROOT}/envs/rlbench"
R3D_VENV="${WORKSPACE_ROOT}/venv_r3d"

echo "[setup] repo_root=${REPO_ROOT}"
echo "[setup] workspace_root=${WORKSPACE_ROOT}"

required_paths=(
  "${WORKSPACE_ROOT}/assets/coppeliasim_v4_1_0"
  "${WORKSPACE_ROOT}/system_shims/nvidia570"
  "${WORKSPACE_ROOT}/models/openai_clip_vit_base_patch32"
  "${WORKSPACE_ROOT}/third_party/RLBench_peract2_20260324"
  "${WORKSPACE_ROOT}/third_party/PyRep_peract2_20260324"
  "${WORKSPACE_ROOT}/third_party/YARR_peract2_20260324"
  "${WORKSPACE_ROOT}/third_party/peract2_bimanual_probe_20260324"
)

for path in "${required_paths[@]}"; do
  if [[ ! -e "${path}" ]]; then
    echo "[setup] missing required path: ${path}" >&2
    exit 1
  fi
done

if [[ ! -x /workspace/.tools/micromamba/bin/micromamba ]]; then
  echo "[setup] missing micromamba at /workspace/.tools/micromamba/bin/micromamba" >&2
  exit 1
fi

if [[ ! -d "${RLBENCH_ENV}" ]]; then
  echo "[setup] creating RLBench env from explicit lock file"
  /workspace/.tools/micromamba/bin/micromamba create -y -p "${RLBENCH_ENV}" --file "${REPO_ROOT}/environment/rlbench_env_explicit.txt"
else
  echo "[setup] RLBench env already exists at ${RLBENCH_ENV}"
fi

if [[ ! -d "${R3D_VENV}" ]]; then
  echo "[setup] creating ${R3D_VENV}"
  python3 -m venv "${R3D_VENV}"
fi

"${R3D_VENV}/bin/pip" install --upgrade pip
"${R3D_VENV}/bin/pip" install -r "${REPO_ROOT}/environment/reveal_pip_freeze.txt"

ln -sfn "${REPO_ROOT}/code/reveal_vla_bimanual" "${WORKSPACE_ROOT}/reveal_vla_bimanual"

echo "[setup] source ${REPO_ROOT}/environment/runtime_env_vars.sh before RLBench commands"
echo "[setup] run ${REPO_ROOT}/environment/validate_same_machine.sh for validation"