alienv

#!/bin/bash -e

# Load and test Modulefiles created by recipes.
PROG=$(basename "$0")

# Detect parent shell (fall back to bash)
[[ -z "$MODULES_SHELL" ]] && MODULES_SHELL=$(ps -e -o pid,command | grep -E "^\s*$PPID\s+" | awk '{print $2}' | sed -e 's/^-\{0,1\}\(.*\)$/\1/')
case "$MODULES_SHELL" in
  sh)                                      ;;
  csh|tcsh) SHELL_NORC_PARAM=-f            ;;
  ksh)      SHELL_NORC_ENV="ENV=/dev/null" ;;
  zsh)      SHELL_NORC_PARAM=--no-rcs      ;;
  *)        MODULES_SHELL=bash
            SHELL_NORC_PARAM=--norc        ;;
esac

function printHelp() {
local EM=`printf '\033[35m'`
local ET=`printf '\033[36m'`
local EZ=`printf '\033[m'`
cat >&2 <<EoF
$PROG -- load environment for aliBuild packages through modulefiles

Usage: $0 \\
         [--architecture|-a ${ET}ARCHITECTURE${EZ}] \\
         [--work-dir|-w ${ET}WORKDIR${EZ}]          \\
         [--no-refresh]                   \\
         ${ET}COMMAND...${EZ}

  ${EM}--no-refresh${EZ} skips refreshing the modules directory.

  ${ET}WORKDIR${EZ} defaults to sw.
  ${ET}ARCHITECTURE${EZ} is automatically detected in most cases.

  ${ET}COMMAND...${EZ} might be:

    ${EM}help${EZ}
      This help screen.

    ${EM}enter${EZ} [${ET}--shellrc${EZ}] MODULE1[,MODULE2...]
      Enters a new shell with the given modules loaded.
      Return to the clean environment by exiting the shell with ${ET}exit${EZ}.
      Inside the environment you can use the native ${ET}modulecmd${EZ}.
      By default you enter the same shell type you are in. Override with environment variable ${ET}MODULES_SHELL${EZ}.
      The new shell will not load your shell startup file by default (e.g. ~/.bashrc) to avoid environment conflicts.
      If you want to retain your shell startup files use ${ET}--shellrc${EZ}.
      Alternatively you can use the ${EM}load${EZ} command in the current shell.

    ${EM}setenv${EZ} MODULE1[,MODULE2...] ${ET}-c${EZ} cmdInEnvironment [PARAM1 [PARAM2...]]
      Executes the given command with environment defined by the given modules.
      Everything after ${ET}-c${EZ} is executed as-is.
      Exit code is preserved.
      Example: $ET$0 setenv AliRoot/v5-08-02-1 -c aliroot -b$EZ

    ${EM}printenv${EZ} or ${EM}load${EZ} [${ET}-q${EZ}] MODULE1[,MODULE2...]
      Prints the environment in the current shell for the given modules (${ET}-q${EZ} for quiet).
      This command does not set any environment and it must be executed through ${ET}eval${EZ} to be effective.
      Override shell with the environment variable ${ET}MODULES_SHELL${EZ}.
      Example: ${ET}eval \`$0 load AliRoot/latest\`$EZ (those are backquotes!)

    ${EM}unload${EZ} [${ET}-q${EZ}] MODULE1[,MODULE2...]
      Prints the environment in the current shell for unloading the given modules (${ET}-q${EZ} for quiet).
      This command does not set any environment and it must be executed through ${ET}eval${EZ} to be effective.
      Override shell with the environment variable ${ET}MODULES_SHELL${EZ}.
      Example: ${ET}eval \`$0 unload AliRoot\`$EZ (version can be omitted)

    ${EM}q${EZ} or ${EM}query${EZ} [REGEXP]
      List all available modules, or the ones matching ${ET}REGEXP${EZ} if provided.

    ${EM}list${EZ}
      List loaded modules.

    ${EM}modulecmd${EZ} [PARAM1 [PARAM2...]]]
      Pass all arguments as-is to the ${ET}modulecmd${EZ} command.
      Example: print AliRoot env for zsh: $ET$0 modulecmd zsh load AliRoot/v5-08-02-1$EZ
      Consult ${ET}man modulecmd${EZ} for more information.

    ${EM}shell-helper${EZ}
      Returns a script to be ${ET}eval${EZ}d in your shell rc file to allow easier loading/unloading.
      In your shell rc put:
        ${ET}ALICE_WORK_DIR=<path_to_alibuild_sw_dir>${EZ}
        ${ET}eval "\`alienv shell-helper\`"${EZ}

EoF
[[ -z "$1" ]] || printf "\033[31mERROR: $1\033[m\n" >&2
}

function installHint() {
  if [[ `uname` == Darwin ]]; then
    CMD='brew install modules'
  elif which apt-get > /dev/null 2>&1; then
    CMD='apt-get install environment-modules'
  elif which yum > /dev/null 2>&1; then
    CMD='yum install environment-modules'
  fi
  printf "\033[31mERROR: Environment Modules was not found on your system.\n" >&2
  if [[ -z "$CMD" ]]; then
    printf "       The package is usually called \033[35menvironment-modules\033[31m.\n" >&2
  else
    printf "       Get it with: \033[35m$CMD\n" >&2
  fi
  printf "\033[m"
}

function normModules() {
  echo "$@" | sed -e 's/,/ /g; s/VO_ALICE@//g; s!::!/!g'
}

function existModules() {
  local MODULE
  for MODULE in "$@"; do
    $MODULECMD bash -t avail 2>&1 | grep -qFx "$MODULE" || \
      { printf "\033[31mERROR: $MODULE was not found\033[m\n" >&2; return 1; }
  done
}

ARCHITECTURE=
WORK_DIR=
NO_REFRESH=
DEFAULT_WORK_DIRS=( sw ../sw )
ARGS=()
COMMAND_IN_ENV=()
CLEAN_ENV=1
VERBOSE=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    --architecture|-a) ARCHITECTURE="$2"; shift 2 ;;
    --work-dir|-w) WORK_DIR="$2"; shift 2 ;;
    --no-refresh) NO_REFRESH=1; shift ;;
    -q) unset VERBOSE; shift ;;
    --shellrc) unset CLEAN_ENV; shift ;;
    --help|help) printHelp; exit 0 ;;
    -c) shift; COMMAND_IN_ENV=("$@"); break ;;
    *) ARGS+=("$1"); shift ;;
  esac
done

ACTION="${ARGS[0]}"
ARGS=("${ARGS[@]:1}")

if [[ "$ACTION" == "shell-helper" ]]; then
cat <<\EOF
# Source this file from your .bashrc, .zshrc or .kshrc.
# Do not forget to export ALICE_WORK_DIR first.
alienv() {
  if test -z "$ALICE_WORK_DIR"; then
    echo "Export ALICE_WORK_DIR to make the alienv shortcut work." >&2
    return 1
  fi
  ALICE_ALIENV__="$ALICE_ALIENV"
  if test -z "$ALICE_ALIENV__"; then
    ALICE_ALIENV__=`unset alienv; export LC_ALL=C; export LANG=C; type -p alienv | sed -e 's/^alienv is //g' 2> /dev/null`
    test -z "$ALICE_ALIENV__" && ALICE_ALIENV__="$ALICE_WORK_DIR/../alibuild/alienv"
  fi
  for A in "$@"; do
    if ! test `echo "$A" | cut -b1` = -; then
      if test "$A" = load || test "$A" = unload; then
        unset A
        eval `"$ALICE_ALIENV__" -w "$ALICE_WORK_DIR" "$@"`
        return $?
      else
        break
      fi
    elif test "$A" = -c; then
      break
    fi
  done
  unset A
  "$ALICE_ALIENV__" -w "$ALICE_WORK_DIR" "$@"
}
EOF
  exit 0
fi

if [[ -z "$WORK_DIR" ]]; then
  for WORK_DIR in "${DEFAULT_WORK_DIRS[@]}"; do
    [[ -d "$WORK_DIR" ]] && break || WORK_DIR=
  done
  [[ -z "$WORK_DIR" ]] && { printHelp "No default work dir (${DEFAULT_WORK_DIRS[*]}) can be accessed"; false; }
fi
[[ ! -d "$WORK_DIR" ]] && { printHelp "Work dir $WORK_DIR cannot be accessed"; false; }
WORK_DIR=$(cd "$WORK_DIR"; pwd)
[[ -z "$ARCHITECTURE" ]] && ARCHITECTURE=$(ls -1t $WORK_DIR | grep -vE '^[A-Z]+$' | head -n1)
[[ -z "$ARCHITECTURE" ]] && { printHelp "Cannot autodetect architecture"; false; }

MODULECMD=$(PATH="$PATH:`brew --prefix modules 2>/dev/null || true`/Modules/bin" which modulecmd || true)
[[ -z "$MODULECMD" ]] && { installHint; false; }

if [[ -d $WORK_DIR/MODULES/$ARCHITECTURE ]]; then
  touch $WORK_DIR/MODULES/$ARCHITECTURE/.testwrite 2> /dev/null || NO_REFRESH=1
  rm -f $WORK_DIR/MODULES/$ARCHITECTURE/.testwrite 2> /dev/null || true
fi

if [[ -z "$NO_REFRESH" ]]; then
  # Collect all modulefiles in one place
  rm -rf $WORK_DIR/MODULES/$ARCHITECTURE
  mkdir -p $WORK_DIR/MODULES/$ARCHITECTURE/BASE
  cat > $WORK_DIR/MODULES/$ARCHITECTURE/BASE/1.0 <<EOF
#%Module1.0
set base_path $WORK_DIR/$ARCHITECTURE
setenv BASEDIR \$base_path
set osname [uname sysname]
set osarchitecture [uname machine]
EOF
  while read PKG; do
    PKGVER=${PKG##*/}
    PKGNAME=${PKG%/*}
    PKGNAME=${PKGNAME##*/}
    [[ ! -e "$PKG/etc/modulefiles/$PKGNAME" ]] && continue
    mkdir -p "$WORK_DIR/MODULES/$ARCHITECTURE/$PKGNAME"
    cp "$PKG/etc/modulefiles/$PKGNAME" "$WORK_DIR/MODULES/$ARCHITECTURE/$PKGNAME/$PKGVER"
  done < <(find $WORK_DIR/$ARCHITECTURE -maxdepth 2 -mindepth 2)
else
  printf "\033[33mWARNING: not updating modulefiles\033[m\n" >&2
fi

export MODULEPATH="$WORK_DIR/MODULES/$ARCHITECTURE:$MODULEPATH"
IGNORE_ERR="Unable to locate a modulefile for 'Toolchain/"
case "$ACTION" in
  enter)
    MODULES=$(normModules "${ARGS[@]}")
    existModules $MODULES
    eval $($MODULECMD bash add $MODULES 2> >(grep -v "$IGNORE_ERR" >&2))
    if [[ ! -z "$CLEAN_ENV" ]]; then
      case $MODULES_SHELL in
        sh|bash) export PS1="[$MODULES]"' \w $> '   ;;
        ksh)     export PS1="[$MODULES]"' $PWD $> ' ;;
        zsh)     export PS1="[$MODULES]"' %~ %#> '  ;;
      esac
    fi
    $MODULECMD bash list
    printf "\033[36mUse \033[35malienv list\033[36m to list loaded modules. " >&2
    printf "Use \033[35mexit\033[36m to exit this environment.\033[m\n" >&2
    exec ${CLEAN_ENV:+env $SHELL_NORC_ENV} $MODULES_SHELL ${CLEAN_ENV:+$SHELL_NORC_PARAM} -i
  ;;
  printenv|load|unload)
    [[ $ACTION == printenv ]] && ACTION=load
    MODULES=$(normModules "${ARGS[@]}")
    [[ $ACTION == load ]] && existModules $MODULES
    [[ $VERBOSE == 1 ]] && printf "\033[36mUse \033[35malienv list\033[36m to list loaded modules.\033[m\n" >&2
    $MODULECMD $MODULES_SHELL $ACTION $MODULES 2> >(grep -v "$IGNORE_ERR" >&2)
    exit 0
  ;;
  setenv)
    [[ -z "${COMMAND_IN_ENV[*]}" ]] && { printHelp "No command specified with -c"; false; }
    MODULES=$(normModules "${ARGS[@]}")
    existModules $MODULES
    eval $($MODULECMD bash add $MODULES 2> >(grep -v "$IGNORE_ERR" >&2))
    exec "${COMMAND_IN_ENV[@]}"
  ;;
  q|query)
    [[ -z "${ARGS[0]}" ]] && SEARCH_CMD=cat || SEARCH_CMD=( grep -iE "${ARGS[0]}" )
    $MODULECMD bash -t avail 2>&1 | grep -E '^[^/]+/[^/]+$' | grep -vE ':$' | \
      "${SEARCH_CMD[@]}" | sed -e 's!^\([^/]*\)/\(.*\)$!VO_ALICE@\1::\2!'
    exit ${PIPESTATUS[3]}  # grep
  ;;
  avail)
    exec $MODULECMD bash avail
  ;;
  list)
    exec $MODULECMD bash list
  ;;
  modulecmd)
    exec $MODULECMD "${ARGS[@]}"
  ;;
  '')
    printHelp "What do you want to do?"
  ;;
  *)
    printHelp "Unknown command: $1"
  ;;
esac
false