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

ROOT="${1:-/workspace}"
MINIFORGE_DIR="${ROOT}/miniforge3"
ENV_PREFIX="${MINIFORGE_DIR}/envs/rlbench"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
EXTERNAL_ROOT="${REPO_ROOT}/external"

echo "[1/6] optional system packages"
echo "Install these on Ubuntu 22.04 if they are missing:"
echo "  sudo apt-get update"
echo "  sudo apt-get install -y git wget curl xvfb git-lfs libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 libxrandr2 libxi6 libxcursor1 libxinerama1 libxxf86vm1 libfontconfig1 libfreetype6"

echo "[2/6] install Miniforge if needed"
if [[ ! -x "${MINIFORGE_DIR}/bin/conda" ]]; then
  wget -O "${ROOT}/Miniforge3-Linux-x86_64.sh" \
    "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh"
  bash "${ROOT}/Miniforge3-Linux-x86_64.sh" -b -p "${MINIFORGE_DIR}"
fi

source "${MINIFORGE_DIR}/etc/profile.d/conda.sh"

echo "[3/6] create or update rlbench environment"
if [[ ! -d "${ENV_PREFIX}" ]]; then
  conda env create -p "${ENV_PREFIX}" -f "${SCRIPT_DIR}/conda_env_rlbench.yml"
else
  conda env update -p "${ENV_PREFIX}" -f "${SCRIPT_DIR}/conda_env_rlbench.yml" --prune
fi
conda activate "${ENV_PREFIX}"

echo "[4/6] install benchmark source trees"
mkdir -p "${ROOT}"
if [[ -d "${EXTERNAL_ROOT}/rlbench" ]]; then
  rsync -a "${EXTERNAL_ROOT}/rlbench/" "${ROOT}/rlbench/"
  rsync -a "${EXTERNAL_ROOT}/pyrep/" "${ROOT}/pyrep/"
  rsync -a "${EXTERNAL_ROOT}/peract_bimanual/" "${ROOT}/peract_bimanual/"
  rsync -a "${EXTERNAL_ROOT}/yarr/" "${ROOT}/yarr/"
else
  if [[ ! -d "${ROOT}/rlbench/.git" ]]; then
    git clone https://github.com/markusgrotz/RLBench.git "${ROOT}/rlbench"
  fi
  if [[ ! -d "${ROOT}/pyrep/.git" ]]; then
    git clone https://github.com/markusgrotz/PyRep.git "${ROOT}/pyrep"
  fi
  if [[ ! -d "${ROOT}/peract_bimanual/.git" ]]; then
    git clone https://github.com/markusgrotz/peract_bimanual.git "${ROOT}/peract_bimanual"
  fi
  if [[ ! -d "${ROOT}/yarr/.git" ]]; then
    git clone https://github.com/markusgrotz/YARR.git "${ROOT}/yarr"
  fi

  git -C "${ROOT}/rlbench" checkout 8af748c51287989294e00c9c670e3330a0e35ed5
  git -C "${ROOT}/pyrep" checkout b8bd1d7a3182adcd570d001649c0849047ebf197
  git -C "${ROOT}/peract_bimanual" checkout bb0232a6ba3fe116566e9568f0c7af980ed6703d
  git -C "${ROOT}/yarr" checkout 6822ff78602c77878b27d4cfe759ce029c67bffb
fi

echo "[5/6] install editable repos used by the rlbench env"
python -m pip install -e "${ROOT}/pyrep"
python -m pip install -e "${ROOT}/rlbench"

echo "[6/6] CoppeliaSim and dataset"
echo "If the repo includes runtime asset archives, extract them first via:"
echo "  ${SCRIPT_DIR}/extract_runtime_assets.sh ${ROOT}"
echo "Provide a CoppeliaSim extraction at ${ROOT}/coppelia_sim."
echo "Provide the benchmark dataset at ${ROOT}/data/bimanual_take_tray_out_of_oven_train_128."
echo "Then source: ${SCRIPT_DIR}/activate_rlbench_runtime.sh ${ROOT}"