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

ROOT_DIR="${ROOT_DIR:-/workspace}"
PROJECT_DIR="${ROOT_DIR}/reveal_vla_bimanual"
MAMBA_BIN="$("${PROJECT_DIR}/scripts/install_micromamba.sh")"
MAMBA_ROOT_PREFIX="${ROOT_DIR}/.micromamba"
ENV_PREFIX="${ROOT_DIR}/envs/rlbench"
ENV_YAML="${PROJECT_DIR}/envs/rlbench310.yaml"
COPPELIA_DIR="$("${PROJECT_DIR}/scripts/download_coppeliasim.sh")"
RUNTIME_DIR="${ROOT_DIR}/runtime"

mkdir -p "${ROOT_DIR}/envs"
mkdir -p "${RUNTIME_DIR}"
chmod 700 "${RUNTIME_DIR}"

run_in_env() {
  local driver_version=""
  local driver_branch=""
  local shim_ld=""
  if command -v nvidia-smi >/dev/null 2>&1; then
    driver_version="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1 || true)"
    driver_branch="${driver_version%%.*}"
    if [[ -n "${driver_branch}" && -d "${ROOT_DIR}/system_shims/nvidia${driver_branch}/usr/lib/x86_64-linux-gnu" ]]; then
      shim_ld="${ROOT_DIR}/system_shims/nvidia${driver_branch}/usr/lib/x86_64-linux-gnu"
    fi
  fi
  env \
    COPPELIASIM_ROOT="${COPPELIA_DIR}" \
    XDG_RUNTIME_DIR="${RUNTIME_DIR}" \
    LD_LIBRARY_PATH="${COPPELIA_DIR}${shim_ld:+:${shim_ld}}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" \
    QT_QPA_PLATFORM_PLUGIN_PATH="${COPPELIA_DIR}" \
    "${MAMBA_BIN}" run -r "${MAMBA_ROOT_PREFIX}" -p "${ENV_PREFIX}" "$@"
}

if [[ -x "${ENV_PREFIX}/bin/python" ]]; then
  "${MAMBA_BIN}" install -y -r "${MAMBA_ROOT_PREFIX}" -p "${ENV_PREFIX}" -f "${ENV_YAML}"
else
  "${MAMBA_BIN}" create -y -r "${MAMBA_ROOT_PREFIX}" -p "${ENV_PREFIX}" -f "${ENV_YAML}"
fi

mkdir -p "${ENV_PREFIX}/etc/conda/activate.d"
cat > "${ENV_PREFIX}/etc/conda/activate.d/coppeliasim.sh" <<EOF
export COPPELIASIM_ROOT=${COPPELIA_DIR}
export XDG_RUNTIME_DIR=${RUNTIME_DIR}
_codex_driver_version="\$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1 2>/dev/null || true)"
_codex_driver_branch="\${_codex_driver_version%%.*}"
_codex_nvidia_ld=
if [[ -n "\${_codex_driver_branch}" && -d ${ROOT_DIR}/system_shims/nvidia\${_codex_driver_branch}/usr/lib/x86_64-linux-gnu ]]; then
  _codex_nvidia_ld=${ROOT_DIR}/system_shims/nvidia\${_codex_driver_branch}/usr/lib/x86_64-linux-gnu
fi
export LD_LIBRARY_PATH=\${COPPELIASIM_ROOT}\${_codex_nvidia_ld:+:\${_codex_nvidia_ld}}\${LD_LIBRARY_PATH:+:\${LD_LIBRARY_PATH}}
export QT_QPA_PLATFORM_PLUGIN_PATH=\${COPPELIASIM_ROOT}
EOF

run_in_env python -m pip install -U pip setuptools wheel
run_in_env python -m pip install --force-reinstall --no-deps numpy==1.26.4 Pillow==12.1.1
run_in_env python -m pip install --no-deps moviepy==2.2.1 timeout-decorator==0.5.0 opencv-python==4.10.0.84 pyquaternion==0.9.9 click-prompt==0.5.1
run_in_env python -m pip install --no-deps poetry-core
run_in_env python -m pip install gymnasium==1.0.0a2 open3d==0.19.0 segment-anything==1.0 transforms3d==0.4.1 openai==0.28.1
first_existing_dir() {
  local candidate=""
  for candidate in "$@"; do
    if [[ -d "${candidate}" ]]; then
      printf '%s\n' "${candidate}"
      return 0
    fi
  done
  printf '%s\n' "$1"
}

PERACT_ROOT="$(first_existing_dir "${ROOT_DIR}/third_party/peract2_bimanual_probe_20260324" "${ROOT_DIR}/third_party/peract_bimanual")"
YARR_ROOT="$(first_existing_dir "${ROOT_DIR}/third_party/YARR_peract2_20260324" "${ROOT_DIR}/third_party/YARR")"
PYREP_ROOT="$(first_existing_dir "${ROOT_DIR}/third_party/PyRep_peract2_20260324" "${ROOT_DIR}/third_party/PyRep")"
RLBENCH_ROOT="$(first_existing_dir "${ROOT_DIR}/third_party/RLBench_peract2_20260324" "${ROOT_DIR}/third_party/RLBench")"
if [[ -f "${PERACT_ROOT}/YARR/setup.py" ]]; then
  YARR_ROOT="${PERACT_ROOT}/YARR"
fi
if [[ -f "${PERACT_ROOT}/PyRep/setup.py" ]]; then
  PYREP_ROOT="${PERACT_ROOT}/PyRep"
fi
if [[ -f "${PERACT_ROOT}/RLBench/setup.py" ]]; then
  RLBENCH_ROOT="${PERACT_ROOT}/RLBench"
fi
if [[ ! -f "${PERACT_ROOT}/pyproject.toml" && ! -f "${PERACT_ROOT}/setup.py" && -f "${PERACT_ROOT}/peract/setup.py" ]]; then
  PERACT_ROOT="${PERACT_ROOT}/peract"
fi
run_in_env python -m pip install --no-build-isolation -e "${PROJECT_DIR}"
run_in_env python -m pip install --no-build-isolation --no-deps -e "${YARR_ROOT}"
run_in_env python -m pip install --no-build-isolation --no-deps -e "${PYREP_ROOT}"
run_in_env python -m pip install --no-build-isolation --no-deps -e "${RLBENCH_ROOT}"
run_in_env python -m pip install --no-build-isolation --no-deps -e "${PERACT_ROOT}"

printf 'RLBench env ready at %s\n' "${ENV_PREFIX}"
printf 'Activate with:\n'
printf '  eval "$(%s shell hook -s bash -r %s)"\n' "${MAMBA_BIN}" "${MAMBA_ROOT_PREFIX}"
printf '  micromamba activate %s\n' "${ENV_PREFIX}"