#!/usr/bin/env bash
# Load HPC-X (module-less) env inside the container
# Usage: source /path/to/env.sh
# Optional: source /path/to/env.sh --unset   # 清理本脚本注入的变量与路径

HPCX_BASE="/data/install/hpcx-v2.21.3-gcc-doca_ofed-ubuntu22.04-cuda12-x86_64"

# --------------- helpers ---------------
_contains_path() { [[ ":$1:" == *":$2:"* ]]; }
prepend_path() {
  local var="$1" val="$2"
  eval "cur=\"\${$var:-}\""
  # 初始化为空时
  if [[ -z "$cur" ]]; then eval "export $var=\"$val\""; return; fi
  # 避免重复
  if _contains_path "$cur" "$val"; then eval "export $var=\"$cur\""; else eval "export $var=\"$val:\$cur\""; fi
}

if [[ "${1-}" == "--unset" ]]; then
  # 基于固定集合，从 PATH/LD_LIBRARY_PATH 等移除相关条目
  PATH_ENTRIES=("$HPCX_BASE/ucx/bin" "$HPCX_BASE/ucc/bin" "$HPCX_BASE/hcoll/bin" "$HPCX_BASE/sharp/bin" \
                "$HPCX_BASE/ompi/tests/imb" "$HPCX_BASE/clusterkit/bin" "$HPCX_BASE/ompi/bin")
  LD_ENTRIES=("$HPCX_BASE/ucx/lib" "$HPCX_BASE/ucx/lib/ucx" "$HPCX_BASE/ucc/lib" "$HPCX_BASE/ucc/lib/ucc" \
              "$HPCX_BASE/hcoll/lib" "$HPCX_BASE/sharp/lib" "$HPCX_BASE/ompi/lib" \
              "$HPCX_BASE/ncclnet_plugin/lib" "$HPCX_BASE/nccl_rdma_sharp_plugin/lib")
  LIB_ENTRIES=("${LD_ENTRIES[@]}")
  CPATH_ENTRIES=("$HPCX_BASE/hcoll/include" "$HPCX_BASE/sharp/include" "$HPCX_BASE/ucx/include" "$HPCX_BASE/ucc/include" "$HPCX_BASE/ompi/include")
  PKG_ENTRIES=("$HPCX_BASE/hcoll/lib/pkgconfig" "$HPCX_BASE/sharp/lib/pkgconfig" "$HPCX_BASE/ucx/lib/pkgconfig" "$HPCX_BASE/ompi/lib/pkgconfig")
  MAN_ENTRIES=("$HPCX_BASE/ompi/share/man")

  clean_var() {
    local var="$1"; shift
    eval "cur=\"\${$var:-}\""
    IFS=':' read -r -a parts <<< "$cur"
    local out=()
    for p in "${parts[@]}"; do
      skip=0; for rm in "$@"; do [[ "$p" == "$rm" ]] && { skip=1; break; }; done
      (( skip==0 )) && out+=("$p")
    done
    local joined="${out[*]}"; joined="${joined// /:}"
    eval "export $var=\"$joined\""
  }

  clean_var PATH "${PATH_ENTRIES[@]}"
  clean_var LD_LIBRARY_PATH "${LD_ENTRIES[@]}"
  clean_var LIBRARY_PATH "${LIB_ENTRIES[@]}"
  clean_var CPATH "${CPATH_ENTRIES[@]}"
  clean_var PKG_CONFIG_PATH "${PKG_ENTRIES[@]}"
  clean_var MANPATH "${MAN_ENTRIES[@]}"

  unset HPCX_DIR HPCX_HOME HPCX_UCX_DIR HPCX_UCC_DIR HPCX_SHARP_DIR HPCX_HCOLL_DIR \
        HPCX_NCCLNET_PLUGIN_DIR HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR HPCX_CLUSTERKIT_DIR \
        HPCX_MPI_DIR HPCX_OSHMEM_DIR HPCX_MPI_TESTS_DIR HPCX_OSU_DIR HPCX_OSU_CUDA_DIR \
        OPAL_PREFIX PMIX_INSTALL_PREFIX OMPI_HOME MPI_HOME OSHMEM_HOME SHMEM_HOME
  echo "HPC-X env cleaned."
  return 0
fi

# --------------- env exports ---------------
export HPCX_DIR="$HPCX_BASE"
export HPCX_HOME="$HPCX_BASE"
export HPCX_UCX_DIR="$HPCX_BASE/ucx"
export HPCX_UCC_DIR="$HPCX_BASE/ucc"
export HPCX_SHARP_DIR="$HPCX_BASE/sharp"
export HPCX_HCOLL_DIR="$HPCX_BASE/hcoll"
export HPCX_NCCLNET_PLUGIN_DIR="$HPCX_BASE/ncclnet_plugin"
export HPCX_NCCL_RDMA_SHARP_PLUGIN_DIR="$HPCX_BASE/nccl_rdma_sharp_plugin"
export HPCX_CLUSTERKIT_DIR="$HPCX_BASE/clusterkit"
export HPCX_MPI_DIR="$HPCX_BASE/ompi"
export HPCX_OSHMEM_DIR="$HPCX_BASE/ompi"
export HPCX_MPI_TESTS_DIR="$HPCX_BASE/ompi/tests"
export HPCX_OSU_DIR="$HPCX_BASE/ompi/tests/osu-micro-benchmarks"
export HPCX_OSU_CUDA_DIR="$HPCX_BASE/ompi/tests/osu-micro-benchmarks-cuda"

prepend_path PATH "$HPCX_BASE/ucx/bin"
prepend_path PATH "$HPCX_BASE/ucc/bin"
prepend_path PATH "$HPCX_BASE/hcoll/bin"
prepend_path PATH "$HPCX_BASE/sharp/bin"
prepend_path PATH "$HPCX_BASE/ompi/tests/imb"
prepend_path PATH "$HPCX_BASE/clusterkit/bin"
prepend_path PATH "$HPCX_BASE/ompi/bin"

prepend_path LD_LIBRARY_PATH "$HPCX_BASE/ucx/lib"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/ucx/lib/ucx"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/ucc/lib"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/ucc/lib/ucc"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/hcoll/lib"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/sharp/lib"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/ompi/lib"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/ncclnet_plugin/lib"
prepend_path LD_LIBRARY_PATH "$HPCX_BASE/nccl_rdma_sharp_plugin/lib"

prepend_path LIBRARY_PATH "$HPCX_BASE/ucx/lib"
prepend_path LIBRARY_PATH "$HPCX_BASE/ucc/lib"
prepend_path LIBRARY_PATH "$HPCX_BASE/hcoll/lib"
prepend_path LIBRARY_PATH "$HPCX_BASE/sharp/lib"
prepend_path LIBRARY_PATH "$HPCX_BASE/ompi/lib"
prepend_path LIBRARY_PATH "$HPCX_BASE/ncclnet_plugin/lib"
prepend_path LIBRARY_PATH "$HPCX_BASE/nccl_rdma_sharp_plugin/lib"

prepend_path CPATH "$HPCX_BASE/hcoll/include"
prepend_path CPATH "$HPCX_BASE/sharp/include"
prepend_path CPATH "$HPCX_BASE/ucx/include"
prepend_path CPATH "$HPCX_BASE/ucc/include"
prepend_path CPATH "$HPCX_BASE/ompi/include"

prepend_path PKG_CONFIG_PATH "$HPCX_BASE/hcoll/lib/pkgconfig"
prepend_path PKG_CONFIG_PATH "$HPCX_BASE/sharp/lib/pkgconfig"
prepend_path PKG_CONFIG_PATH "$HPCX_BASE/ucx/lib/pkgconfig"
prepend_path PKG_CONFIG_PATH "$HPCX_BASE/ompi/lib/pkgconfig"

prepend_path MANPATH "$HPCX_BASE/ompi/share/man"

export OPAL_PREFIX="$HPCX_BASE/ompi"
export PMIX_INSTALL_PREFIX="$HPCX_BASE/ompi"
export OMPI_HOME="$HPCX_BASE/ompi"
export MPI_HOME="$HPCX_BASE/ompi"
export OSHMEM_HOME="$HPCX_BASE/ompi"
export SHMEM_HOME="$HPCX_BASE/ompi"

echo "HPC-X env loaded from: $HPCX_BASE"