#!/bin/bash -e

BITSDIR=`dirname $0`

ARGV=("$@"); ARGC=("$#")

function printHelp() {
cat >&2 <<EoF
$PROG -- load environment for bits 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}] [${ET}--dev${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}.
      If you want to add environment variables for development use ${ET}--dev${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}BITS_WORK_DIR=<path_to_bits_sw_dir>${EZ}
        ${ET}eval "\`bits shell-helper\`"${EZ}

EoF
[[ -z "$1" ]] || printf "${ER}ERROR: $*${EZ}\n" >&2
}

function installHint() {
  if [[ $UNAME == Darwin ]]; then
    CMD='brew install modules'
  elif command -v apt-get > /dev/null 2>&1; then
    CMD='apt-get install environment-modules'
  elif command -v yum > /dev/null 2>&1; then
    CMD='yum install environment-modules'
  fi
  printf "${ER}ERROR: Environment Modules was not found on your system.\n" >&2
  if [[ -z "$CMD" ]]; then
    printf "       The package is usually called ${EM}environment-modules${ER}.\n" >&2
  else
    printf "       Get it with: ${EM}${CMD}\n" >&2
  fi
  printf "${EZ}"
}

function detectShell() {
  # 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 collectModules() {
  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 2> /dev/null)
  else
    printf "${EY}WARNING: not updating modulefiles${EZ}\n" >&2
  fi
}

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

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

function stripDyld() {
  TO_STRIP=$( (typeset | grep ^DYLD | cut -d= -f1 | grep '_modshare$' | xargs echo ) 2> /dev/null )
  [ ! "X$TO_STRIP" = X ] && echo unset $TO_STRIP && echo ';'
}

function readBitsRc() {
  cfgfile=$1
  tmpfile=`mktemp`
  if [ -f $cfgfile ]
  then   
    awk '/^\[bits]$/,/^$/ {print}' $cfgfile | sed 's/ *= */=/g' | grep  = > $tmpfile &&  . $tmpfile
    awk "/^\[$organisation]$/,/^$/ {print}" $cfgfile | sed 's/ *= */=/g' | grep -v "^\[$organisation\]" | grep = > $tmpfile && . $tmpfile
  fi
  rm -rf $tmpfile
}

function configBits() {

	# Backward compatibility defaults

	BITS_BRANDING=${BITS_BRANDING:-"bits"}
	BITS_PATH=${BITS_PATH:-""}
	BITS_REPO_DIR=${BITS_REPO_DIR:-"alidist"}
	BITS_WORK_DIR=${BITS_WORK_DIR:-"sw"}
	BITS_ORGANISATION=${BITS_ORGANISATION:-"ALICE"}
	BITS_PKG_PREFIX=${BITS_PKG_PREFIX:-"VO_"$BITS_ORGANISATION}
    
    cfile=""

    for cfg in  $1 bits.rc .bitsrc $HOME/.bitsrc 
    do
	  if [ "x$cfg" !=  "x" -a -f $cfg ]
	  then
	    cfile=$cfg
	    break
      fi
    done
    
    if [ "x$cfile" != "x" ] 
    then
	  readBitsRc $cfile
    fi
    
	export BITS_ORGANISATION=${organisation:-$BITS_ORGANISATION}
	export BITS_BRANDING=${branding:-$BITS_BRANDING}
	export BITS_PATH=${search_path:-$BITS_PATH}
	export BITS_REPO_DIR=${repo_dir:-$BITS_REPO_DIR}
	export BITS_WORK_DIR=${sw_dir:-$BITS_WORK_DIR}
	export BITS_PKG_PREFIX=${pkg_prefix:-$BITS_PKG_PREFIX}

	mkdir -p $BITS_WORK_DIR || echo "Cannot create directory: " BITS_WORK_DIR 
}

config_file=""

for ((i=0;i<$ARGC;i++)); do
    arg=${ARGV[$i]}
    opt=${arg%%=*}
    val=${arg##*=}
    case $opt in
       --config)
          config_file=$val
       ;;
       *)
         new_args+=("$arg")
         ;;
   esac
done

configBits "$config_file"
set -- "${new_args[@]}"

for arg in $@
do
  case $arg in
    analytics|architecture|build|clean|deps|doctor|init|version)
        $BITSDIR/bitsBuild $@
	    exit $?
        ;;
    -h|--help|-d|--debug|-n|--dry-run)
        configBits $config_file
        $BITSDIR/bitsBuild $@
	    exit $?
        ;;
    setenv)
	break
	;;
	 *)
	    shift
        ;;
  esac
done

# If we are on a TTY, enable colors
if test -t 1; then
  EM=`printf '\033[35m'`  # magenta
  ET=`printf '\033[36m'`  # cyan
  ER=`printf '\033[31m'`  # red
  EY=`printf '\033[33m'`  # yellow
  EZ=`printf '\033[m'`    # reset
else
  EM=
  ET=
  ER=
  EY=
  EZ=
fi
  
# My operating system

UNAME=$(uname)

ARCHITECTURE=
WORK_DIR=
NO_REFRESH=
DEFAULT_WORK_DIRS=("$BITS_WORK_DIR" "$BITS_WORK_DIR" ${BITSBUILD_CHDIR:+"$BITSBUILD_CHDIR/sw"} sw ../sw)
COMMAND_IN_ENV=()
CLEAN_ENV=1
DEVOPT=0
VERBOSE=1

ARGS=()

for ((i=0;i<$ARGC;i++));
do
  case "${ARGV[$i]}" in
      --architecture|-a)
	  i+=1 ;ARCHITECTURE="${ARGV[$i]}";;
      --work-dir|-w)
	  i+=1; WORK_DIR="${ARGV[$i]}";;
      --no-refresh)
	  NO_REFRESH=1;;
      -q)
	  unset VERBOSE;;
      --shellrc)
	  unset CLEAN_ENV;;
      --dev)
	  DEVOPT=1;;
      --help|help)
	  printHelp; exit 0 ;;
      -c) 
	  i=$((i+1)); COMMAND_IN_ENV=("${ARGV[@]:${i}}");
	  break ;;
      *)
	  ARGS+=(${ARGV[$i]});;
   esac
done

ACTION="${ARGS[0]}"

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

export BITSLVL=$((BITSLVL+1))

if [[ "$ACTION" == "shell-helper" ]]; then
cat <<\EOF
# Source this file from your .bashrc, .zshrc or .kshrc.
BITS_ENV__="$BITS_ENV"
if test -z "$BITS_ENV__"; then
  BITS_ENV__=`unset bits; export LC_ALL=C; export LANG=C; type -p bits | sed -e 's/^bits is //g' 2> /dev/null`
fi
bits() {
  if test -z "$BITS_ENV__"; then
    echo "bits not found"
    return 1
  fi
  for A in "$@"; do
    if ! test `echo "$A" | cut -b1` = -; then
      if test "$A" = load || test "$A" = unload; then
        unset A
        eval `"$BITS_ENV__" "$@"`
        return $?
      else
        break
      fi
    elif test "$A" = -c; then
      break
    fi
  done
  unset A
  "$BITS_ENV__" "$@"
}
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 can be accessed:"                      \
                                      "export \$BITS_WORK_DIR or run bits from a directory" \
                                      "containing the \"sw\" dir"; false; }
fi
[[ ! -d "$WORK_DIR" ]] && { printHelp "Work dir $WORK_DIR cannot be accessed"; false; }
WORK_DIR=$(cd "$WORK_DIR"; pwd)
[[ -z "$ARCHITECTURE" ]] && ARCHITECTURE="$("bitsBuild" architecture 2> /dev/null || true)"
[[ -z "$ARCHITECTURE" ]] && ARCHITECTURE="$("$BITSDIR/bitsBuild" architecture 2> /dev/null || true)"
[[ -z "$ARCHITECTURE" || "$ARCHITECTURE" == "<unknown>" ]] && ARCHITECTURE=$(ls -1t $WORK_DIR | grep -vE '^[A-Z]+$' | head -n1)
[[ -z "$ARCHITECTURE" ]] && { printHelp "Cannot autodetect architecture"; false; }

# Look for modulecmd (v3) or modulecmd-compat (>= v4)
MODULECMD=$(command -v modulecmd 2> /dev/null || true)
[[ -x "$MODULECMD" ]] || MODULECMD="$(dirname $(command -v envml 2> /dev/null || true) 2> /dev/null || true)/../libexec/modulecmd-compat"
[[ -x "$MODULECMD" ]] || MODULECMD="$(brew --prefix modules 2> /dev/null || true)/libexec/modulecmd-compat"
[[ -x "$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

export MODULEPATH="$WORK_DIR/MODULES/$ARCHITECTURE${MODULEPATH:+":$MODULEPATH"}"
MODULEPATH=$(echo "$MODULEPATH" | sed -e 's/::*/:/g; s/^://; s/:$//')
IGNORE_ERR="Unable to locate a modulefile for 'Toolchain/"
case "$ACTION" in
  enter)
    [[ $BITSLVL == 1 ]] || \
      { printf "${ER}ERROR: already in an bits environment${EZ}\n" >&2; exit 1; }
    MODULES=$(normModules "${ARGS[@]}")
    collectModules
    existModules $MODULES
    PS1DEV=
    if [[ $DEVOPT == 1 ]];then
      PS1DEV=" (dev)"
      for MODULE in $MODULES; do
        . "$WORK_DIR/$ARCHITECTURE/$MODULE/etc/profile.d/init.sh"
      done
    else
      eval $($MODULECMD bash add $MODULES 2> >(grep -v "$IGNORE_ERR" >&2))
    fi
    [[ $UNAME == Darwin ]] && eval $(stripDyld)
    detectShell
    if [[ ! -z "$CLEAN_ENV" ]]; then
      case $MODULES_SHELL in
        sh|bash) export PS1="[$MODULES]$PS1DEV"' \w $> '   ;;
        ksh)     export PS1="[$MODULES]$PS1DEV"' $PWD $> ' ;;
        zsh)     export PS1="[$MODULES]$PS1DEV"' %~ %#> '  ;;
      esac
    fi
    $MODULECMD bash list
    if [[ $DEVOPT == 1 ]];then
      echo "*** Development env loaded for '$MODULES' ***" >&2
    else
      printf "${ET}Use ${EM}bits list${ET} to list loaded modules. " >&2
    fi
    printf "Use ${EM}exit${ET} to exit this environment.${EZ}\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 ]] && { collectModules; existModules $MODULES; }
    [[ $VERBOSE == 1 ]] && printf "${ET}Use ${EM}bits list${ET} to list loaded modules.${EZ}\n" >&2
    if [[ $DEVOPT == 1 ]];then
      echo "Warning: Development mode for '$ACTION' is not supported via bits. Please remove '--dev' or run manually shown below" >&2
      echo "*** Please run following commands manually to load dev environment for '$MODULES' ***" >&2
      echo "" >&2
      echo "export WORK_DIR=$WORK_DIR" >&2
      for MODULE in $MODULES; do
        echo ". $WORK_DIR/$ARCHITECTURE/$MODULE/etc/profile.d/init.sh" >&2
      done
      echo "" >&2
      exit 0
    fi
    detectShell
    $MODULECMD $MODULES_SHELL $ACTION $MODULES 2> >(grep -v "$IGNORE_ERR" >&2)
    [[ $UNAME == Darwin ]] && ( eval $($MODULECMD $MODULES_SHELL $ACTION $MODULES 2> /dev/null) &> /dev/null; stripDyld )
    exit 0
  ;;
  setenv)
    [[ $DEVOPT == 1 ]] && echo "Warning: Development mode for '$ACTION' is not supported. Please remove '--dev'" >&2
    [[ -z "${COMMAND_IN_ENV[*]}" ]] && { printHelp "No command specified with -c"; false; }
    MODULES=$(normModules "${ARGS[@]}")
    collectModules
    existModules $MODULES
    eval $($MODULECMD bash add $MODULES 2> >(grep -v "$IGNORE_ERR" >&2))
    [[ $UNAME == Darwin ]] && eval $(stripDyld)
    exec "${COMMAND_IN_ENV[@]}"
  ;;
  q|query)
    [[ -z "${ARGS[0]}" ]] && SEARCH_CMD=cat || SEARCH_CMD=( grep -iE "${ARGS[0]}" )
    collectModules
    $MODULECMD bash -t avail 2>&1 | grep -E '^[^/]+/[^/]+$' | grep -vE ':$' | \
      "${SEARCH_CMD[@]}" | sed -e 's!^\([^/]*\)/\(.*\)$!'${BITS_PKG_PREFIX}'@\1::\2!'
    exit ${PIPESTATUS[3]}  # grep
  ;;
  avail)
    collectModules
    exec $MODULECMD bash avail
  ;;
  list)
    exec $MODULECMD bash list
  ;;
  modulecmd)
    collectModules
    exec $MODULECMD "${ARGS[@]}"
    ;;
  '')
    $BITSDIR/bitsBuild --help
    printHelp "What do you want to do?"
    ;;
  *)
    printHelp "Unknown command: $1"
  ;;
esac

false
