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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
TASK_ENV="${SEG_TASK_CONDA_ENV:-seg_smp}"
MMSEG_ENV="${SEG_MMSEG_CONDA_ENV:-seg_mmcv}"

env_exists() {
  conda env list | awk '{print $1}' | grep -Fxq "$1"
}

create_task_env() {
  if ! env_exists "${TASK_ENV}"; then
    conda create -n "${TASK_ENV}" python=3.11 -y
  fi
  conda run -n "${TASK_ENV}" python -m pip install -U pip
  conda run -n "${TASK_ENV}" python -m pip install -r "${ROOT_DIR}/backend/requirements.txt"
  conda run -n "${TASK_ENV}" python -m pip install \
    torch==2.6.0 torchvision==0.21.0 \
    'numpy<2' 'opencv-python<4.12' albumentations segmentation-models-pytorch ultralytics \
    mmengine mmsegmentation==1.2.2 mmcv-lite \
    matplotlib pandas scikit-learn scipy seaborn tqdm tensorboard
}

create_mmseg_env() {
  if ! env_exists "${MMSEG_ENV}"; then
    conda create -n "${MMSEG_ENV}" python=3.10 -y
  fi
  conda run -n "${MMSEG_ENV}" python -m pip install -U pip
  conda run -n "${MMSEG_ENV}" python -m pip install \
    torch==2.1.2 torchvision==0.16.2 \
    --index-url https://download.pytorch.org/whl/cu121
  conda run -n "${MMSEG_ENV}" python -m pip install \
    mmengine==0.10.7 mmsegmentation==1.2.2 'mmcv==2.1.0' \
    -f https://download.openmmlab.com/mmcv/dist/cu121/torch2.1/index.html
  conda run -n "${MMSEG_ENV}" python -m pip install \
    'numpy<2' 'opencv-python<4.12' ftfy regex matplotlib pandas scikit-learn scipy seaborn tqdm tensorboard
}

case "${1:-all}" in
  all)
    create_task_env
    create_mmseg_env
    PYTHONPATH="${ROOT_DIR}/backend" conda run -n "${TASK_ENV}" python "${ROOT_DIR}/scripts/verify_runtime_envs.py" --refresh
    ;;
  task)
    create_task_env
    echo "Created or repaired ${TASK_ENV}. Run '$0 all' for full runtime verification."
    ;;
  mmseg)
    create_mmseg_env
    echo "Created or repaired ${MMSEG_ENV}. Run '$0 all' for full runtime verification."
    ;;
  *)
    echo "usage: $0 [all|task|mmseg]" >&2
    exit 2
    ;;
esac