#!/usr/bin/env bash
# One-time (idempotent) setup for sim-office-nav-sem:
#   conda ros2 + ML deps + colcon build with conda python + entry-point shebang fix.
#
# Usage (in card terminal):
#   conda activate ros2
#   bash /workspace/bootstrap_ros2_conda_workspace.sh
set -euo pipefail

ROOT="${OFFICE_VLM_ROOT:-/workspace}"
CONDA_ENV="${ROS2_CONDA_ENV:-ros2}"
MARKER="${ROOT}/.ros2_conda_workspace_ready"
CLIP_VENDOR="${ROOT}/vendor/CLIP"
CLIP_COMMIT="${CLIP_GIT_COMMIT:-d05afc436d78f1c48dc0dbf8e5980a9d471f35f6}"

log() { echo "[bootstrap-ros2] $*"; }

if [[ ! -f "${HOME}/miniconda3/etc/profile.d/conda.sh" ]]; then
  log "ERROR: miniconda3 not found under ${HOME}; run prod_juicefs_clone_ros2_env.sh on cluster first"
  exit 1
fi

# shellcheck source=/dev/null
source "${HOME}/miniconda3/etc/profile.d/conda.sh"
conda activate "${CONDA_ENV}"

export PYTHONNOUSERSITE=1
PY="${CONDA_PREFIX}/bin/python3"
PIP="${CONDA_PREFIX}/bin/pip"

log "conda env=${CONDA_ENV} python=$(${PY} --version 2>&1)"

# Remove wrong PyPI "clip" (clipboard CLI) and ~/.local pollution from old docs.
"${PIP}" uninstall -y clip open-clip-torch 2>/dev/null || true
rm -rf "${HOME}/.local/lib/python3.10/site-packages/clip"* 2>/dev/null || true

log "installing ML + colcon build deps into conda (not --user)"
"${PIP}" install --upgrade pip
"${PIP}" install numpy==1.26.4 \
  torch==2.6.0 torchvision==0.21.0 \
  "opencv-python<4.11" Pillow ftfy regex tqdm matplotlib \
  catkin_pkg empy lark-parser

_ensure_openai_clip() {
  if "${PY}" -c "import clip; assert hasattr(clip, 'load')" 2>/dev/null; then
    log "openai CLIP already available"
    return 0
  fi
  mkdir -p "$(dirname "${CLIP_VENDOR}")"
  touch "$(dirname "${CLIP_VENDOR}")/COLCON_IGNORE"
  if [[ ! -f "${CLIP_VENDOR}/clip/clip.py" ]]; then
    log "cloning OpenAI CLIP -> ${CLIP_VENDOR}"
    rm -rf "${CLIP_VENDOR}"
    for url in \
      "https://ghproxy.net/https://github.com/openai/CLIP.git" \
      "https://github.com/openai/CLIP.git"; do
      if git clone --depth 1 "${url}" "${CLIP_VENDOR}" 2>/dev/null; then
        break
      fi
      rm -rf "${CLIP_VENDOR}"
    done
    [[ -f "${CLIP_VENDOR}/clip/clip.py" ]] || {
      log "ERROR: failed to fetch OpenAI CLIP source"
      exit 1
    }
    if [[ -n "${CLIP_COMMIT}" ]]; then
      git -C "${CLIP_VENDOR}" fetch --depth 1 origin "${CLIP_COMMIT}" 2>/dev/null \
        && git -C "${CLIP_VENDOR}" checkout "${CLIP_COMMIT}" 2>/dev/null || true
    fi
  fi
  "${PIP}" install "${CLIP_VENDOR}"
  "${PY}" -c "import clip; assert hasattr(clip, 'load')"
  log "openai CLIP installed"
}

_ensure_openai_clip

set +u
# shellcheck source=/dev/null
source /opt/ros/humble/setup.bash
set -u

cd "${ROOT}"
if [[ -f "${MARKER}" ]] && grep -q "${CONDA_PREFIX}/bin/python3" "${MARKER}" 2>/dev/null; then
  log "workspace already built for this conda python (marker present); rebuilding vlm packages only"
  colcon build --symlink-install \
    --packages-select vlm_perception_pkg vlm_nav_pkg \
    --cmake-args "-DPython3_EXECUTABLE=${PY}"
else
  log "colcon build (full workspace, conda python)"
  colcon build --symlink-install --cmake-args "-DPython3_EXECUTABLE=${PY}" || {
    log "WARN: full build had failures; ensuring vlm packages at minimum"
    colcon build --symlink-install \
      --packages-select vlm_perception_pkg vlm_nav_pkg \
      --cmake-args "-DPython3_EXECUTABLE=${PY}"
  }
fi

_patch_entry_shebangs() {
  local f
  while IFS= read -r -d '' f; do
    if head -1 "$f" | grep -q '^#!/usr/bin/python3'; then
      sed -i "1s|#!/usr/bin/python3|#!${PY}|" "$f"
      log "patched shebang: $f"
    fi
  done < <(find "${ROOT}/install" -type f -perm /111 -print0 2>/dev/null || true)
}

_patch_entry_shebangs

date -Iseconds > "${MARKER}"
echo "CONDA_PREFIX=${CONDA_PREFIX}" >> "${MARKER}"

set +u
# shellcheck source=/dev/null
source "${ROOT}/install/setup.bash"
set -u
"${PY}" -c "import torch, clip, rclpy; print('verify OK', torch.__version__, torch.cuda.is_available())"

log "done. In new terminals run:"
log "  conda activate ${CONDA_ENV}"
log "  export PYTHONNOUSERSITE=1"
log "  source /opt/ros/humble/setup.bash"
log "  source ${ROOT}/install/setup.bash"