setup

#!/usr/bin/env bash

# shellcheck disable=1090,1091,2046,2068,2086,2206

set -aeuo pipefail

INSTALLDIR=$(readlink -f ${INSTALLDIR:-$PWD/conda})

# In all functions, make variables local unless they are either intended for global use, or set in
# a subshell that exits at the end of the function.

cachedir() {
  (
    env_activate base
    echo $CONDA_PREFIX/cache
  )
}

clean_conda_pkgs_maybe() {
  if [[ -n ${EAGLE_CLEAN_CONDA_PKGS+x} ]]; then
    (
      env_activate base
      msg Cleaning conda package cache
      mamba clean --all --yes
    )
  fi
}

cuda_command_file() {
  local path val
  val=${ARGS[cudascript]}
  path=$val
  test -e $path || path=$(dirname $(readlink -f $0))/cuda/$val
  echo $path
}

cuda_release() {
  (
    source $(cuda_command_file)
    nvcc --version | grep release | sed -E 's/^.*release ([0-9]+).*/\1/'
  )
}

env_activate() {
  local env
  env=$1
  source $INSTALLDIR/etc/profile.d/mamba.sh
  # Activation scripts shipped with conda packages may contain statements either expanding undefined
  # variables or exiting with error status, so temporarily disable the checks for these conditions,
  # then enable again.
  set +eu
  mamba activate $env
  set -eu
}

env_exists() {
  local d name x
  name=$1
  d=$INSTALLDIR/envs
  mkdir -p $d
  for x in $(find $d -maxdepth 1 -mindepth 1 -type d | sort); do
    if [[ $(basename $x) == "$name" ]]; then
      msg Found existing conda environment: $name
      return 0
    fi
  done
  return 1
}

install_conda() {
  local installer url ver
  if [[ -d $INSTALLDIR ]]; then
    msg Found existing conda installation: $INSTALLDIR
    return
  fi
  ver=26.1.1-3
  installer=Miniforge3-$ver-$(uname -s)-$(uname -p).sh
  url=https://github.com/conda-forge/miniforge/releases/download/$ver/$installer
  msg Installing conda
  wget -nv $url
  bash $installer -bfp $INSTALLDIR
  rm -v $installer
}

install_devpkgs_maybe() {
  local name
  if [[ -v EAGLEDEV ]]; then
    name=$1
    msg Installing dev packages into environment: $name
    (
      env_activate $name
      mamba env update ${MAMBAFLAGS[@]} --file envs/dev.yaml
    )
  fi
}

msg() {
  echo -e "=> $*"
}

parse_kvargs() {
  local arg key required val
  declare -g -A ARGS
  for arg in $@; do
    IFS="=" read -r key val <<<$arg
    ARGS[$key]=$val
  done
  required="cudascript"
  for arg in $required; do
    test -z "${ARGS[$arg]:-}" && echo Argument $arg= missing && exit 1
  done
  export ARGS
}

prep_env() {
  local args name
  name=$1
  args=(${MAMBAFLAGS[@]} --file envs/$name.yaml --name $name)
  (
    env_activate base
    if env_exists $name; then
      msg Updating environment: $name
      mamba env update ${args[@]} | { grep --line-buffered -v "Requirement already satisfied" || true; }
    else
      msg Creating environment: $name
      export PIP_VERBOSE=1
      mamba env create ${args[@]}
    fi
  )
  cat <<EOF

Note: Apparent pip errors, if any, about not installed or incompatible packages are generally safe to ignore.

EOF
  install_devpkgs_maybe $name
  # When runtime virtual environments are activated, set XDG_CACHE_HOME so that pip caches package
  # files in the runtime area, rather than in the user's home directory, which might impact their
  # quota. When deactivated, environments should reset or unset this value, as appropriate.
  #
  # Always run this step to re-create these scripts, as downstream steps will *append* to them.
  (
    env_activate $name
    dir_activate=$CONDA_PREFIX/etc/conda/activate.d
    mkdir -pv $dir_activate
    cat <<EOF >$dir_activate/eagle.sh
if [[ -v XDG_CACHE_HOME ]]; then
  export EAGLE_OLD_XDG_CACHE_HOME=\$XDG_CACHE_HOME
fi
export XDG_CACHE_HOME=$(cachedir)
EOF
    dir_deactivate=$CONDA_PREFIX/etc/conda/deactivate.d
    mkdir -pv $dir_deactivate
    cat <<EOF >$dir_deactivate/eagle.sh
unset XDG_CACHE_HOME
if [[ -v EAGLE_OLD_XDG_CACHE_HOME ]]; then
  export XDG_CACHE_HOME=\$EAGLE_OLD_XDG_CACHE_HOME
  unset EAGLE_OLD_XDG_CACHE_HOME
fi
EOF
  )
  clean_conda_pkgs_maybe
}

prep_env_anemoi() {
  local name
  name=anemoi
  prep_env $name
  (
    env_activate $name
    # NB: Append to the existing script:
    cat <<EOF >>$CONDA_PREFIX/etc/conda/activate.d/eagle.sh
modules="\$(module --terse list 2>/dev/null | grep -v 'No modules loaded' | tr '\n' ' ')"
if [[ -n "\$modules" ]]; then
  export EAGLE_OLD_MODULES="\$modules"
fi
$(cat $(cuda_command_file))
EOF
    # NB: Append to the existing script:
    cat <<EOF >>$CONDA_PREFIX/etc/conda/deactivate.d/eagle.sh
module purge 2>/dev/null
if [[ -v EAGLE_OLD_MODULES ]]; then
  module load \$EAGLE_OLD_MODULES
  unset EAGLE_OLD_MODULES
fi
EOF
  )
}

prep_env_base() {
  local name
  name=base
  (
    env_activate $name
    msg Updating environment: $name
    mamba env update ${MAMBAFLAGS[@]} --file envs/$name.yaml
    install_devpkgs_maybe $name
    clean_conda_pkgs_maybe
  )
}

prep_env_data() {
  local name
  name=data
  prep_env $name
  # Add logic to the [de]activation scripts to silence run-time complaints and avoid
  # segfaults when importing mpi4py in environments without InfiniBand hardware:
  (
    env_activate $name
    dir_activate=$CONDA_PREFIX/etc/conda/activate.d
    mkdir -pv $dir_activate
    # NB: Append to the existing script:
    cat <<EOF >>$CONDA_PREFIX/etc/conda/activate.d/eagle.sh
if [[ -v FI_PSM3_UUID ]]; then
  export EAGLE_OLD_FI_PSM3_UUID=\$FI_PSM3_UUID
fi
export FI_PSM3_UUID=eagle
if [[ -v FI_PROVIDER ]]; then
  export EAGLE_OLD_FI_PROVIDER=\$FI_PROVIDER
fi
export FI_PROVIDER=tcp
EOF
    dir_deactivate=$CONDA_PREFIX/etc/conda/deactivate.d
    mkdir -pv $dir_deactivate
    # NB: Append to the existing script:
    cat <<EOF >>$dir_deactivate/eagle.sh
unset FI_PSM3_UUID
if [[ -v EAGLE_OLD_FI_PSM3_UUID ]]; then
  export FI_PSM3_UUID=\$EAGLE_OLD_FI_PSM3_UUID
  unset EAGLE_OLD_FI_PSM3_UUID
fi
unset FI_PROVIDER
if [[ -v EAGLE_OLD_FI_PROVIDER ]]; then
  export FI_PROVIDER=\$EAGLE_OLD_FI_PROVIDER
  unset EAGLE_OLD_FI_PROVIDER
fi
EOF
  )
}

prep_env_visualization() {
  prep_env visualization
}

prep_env_wxvx() {
  prep_env wxvx
}

parse_kvargs $@
install_conda
CONDA_OVERRIDE_CUDA=$(cuda_release)
MAMBAFLAGS=(--log-level error --yes)
MAMBA_ROOT_PREFIX=$INSTALLDIR
XDG_CACHE_HOME=$(cachedir)
export CONDA_OVERRIDE_CUDA MAMBA_ROOT_PREFIX XDG_CACHE_HOME
environments=(anemoi base data visualization wxvx)
for name in ${environments[@]}; do
  time prep_env_$name
done
msg Done