#!/bin/bash
# -*- shell-script -*-
##
#
#  Script to update Lmod cache files.
#
#  author: Kenneth Hoste (kenneth.hoste@ugent.be)
#  modified by Robert McLay (mclay@tacc.utexas.edu)
#
#  This script is licensed under the terms of the MIT license reproduced below.
#  This means that Lmod is free software and can be used for both academic
#  and commercial purposes at absolutely no cost.
#
##
#
#  Copyright (C) 2015-2015 Ghent University
#
#  Permission is hereby granted, free of charge, to any person obtaining
#  a copy of this software and associated documentation files (the
#  "Software"), to deal in the Software without restriction, including
#  without limitation the rights to use, copy, modify, merge, publish,
#  distribute, sublicense, and/or sell copies of the Software, and to
#  permit persons to whom the Software is furnished to do so, subject
#  to the following conditions:
#
#  The above copyright notice and this permission notice shall be
#  included in all copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#  NONINFRINGEMENT.  IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
#  BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
#  ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
#  CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#  THE SOFTWARE.
#
##

#
# loosely based on https://github.com/TACC/Lmod/tree/main/contrib/BuildSystemCacheFile/createSystemCacheFile.sh

#
# Script parameters (these can be tweaked via the provided command line options)
#

# exit on any error
set -e

# print debug info (default: False)
DEBUG=0

# Base module path
BASE_MODULEPATH=

# Lmod spider debugging
LMOD_DEBUG=

# location of cache directory
LMOD_CACHE_DIR=

# location of cache timestamp file
LMOD_CACHE_TIMESTAMP_FILE=

# contents for cache timestamp file
LMOD_CACHE_TIMESTAMP_FILE_TXT=

# update reverseMapT cache file
UPDATE_REVERSEMAPT_CACHE=0

UPDATE_JSON_RMAPT=0
UPDATE_XALT_JSON_RMAPT=0

# find PYTHON executable
PYTHON=

#
# Utility functions
#

# print debug info
debug() {
    if [ $DEBUG -ne 0 ]; then
        echo "[DEBUG] $1"
    fi
}

# print error to stderr and exit
error() {
    echo "ERROR: $1" >&2
    exit 1
}

# print warning to stderr
warning() {
    echo "WARNING: $1" >&2
}


# print help w.r.t. script usage and command line options
print_help() {
    echo
    echo "Script to update Lmod cache files"
    echo
    echo "$0 [-options] MODULEPATH"
    echo
    echo "Available command line options:"
    echo
    echo "  -d <cacheDir>"
    echo "      location of Lmod cache directory (default: determine via 'ml --config')"
    echo "  -D"
    echo "      enable debug printing"
    echo "  -h/-H"
    echo "      print (this) help and exit"
    echo "  -t <timestamp.txt>"
    echo "      location of Lmod cache timestamp file (default: determine via 'ml --config')"
    echo "  -r"
    echo "      enable updating reverseMapT cache file (default: only spiderT and dbT cache files)"
    echo "  -R <reverseMapD>"
    echo "      enable updating reverseMapT cache file (default: only spiderT and dbT cache files)"
    echo "  -J"
    echo "      enable update jsonReverseMapT cache file"
    echo "  -K"
    echo "      enable update xalt_rmapT cache file"
    echo "  -X"
    echo "      enable Lmod debugging of spider command"
}

# parse command line options using getopts
# see https://www.tldp.org/LDP/abs/html/internal.html#EX33
parse_cmdline(){
    debug "Parsing command line options..."
    local OPTIND
    while getopts :d:DhmHt:rR:XJK opt; do
        case $opt in
            d)
                LMOD_CACHE_DIR=$OPTARG
                ;;
            D)
                DEBUG=1
                ;;
            h|H)
                print_help
                exit 0
                ;;
            t)
                LMOD_CACHE_TIMESTAMP_FILE=$OPTARG;
                ;;
            r)
                UPDATE_REVERSEMAPT_CACHE=1
                ;;
            R)
                UPDATE_REVERSEMAPT_CACHE=1
		LMOD_REVERSEMAPT_DIR=$OPTARG
                ;;
            X)
                LMOD_DEBUG=1
                ;;
            J)
                UPDATE_JSON_RMAPT=1
                ;;
	    K)
                UPDATE_XALT_JSON_RMAPT=1
                ;;
            \?)
                print_help
                error "Unknown option specified: -$OPTARG"
                ;;
            :)
                print_help
                error "Option -$OPTARG requires an argument"
                ;;
        esac
    done

    # shift command line options so other arguments are accessible via $@ or $1, etc.
    shift $((OPTIND - 1))

    BASE_MODULEPATH=$1
}


# set default values for undefined parameters
set_defaults() {

  if [ -z "${LMOD_CACHE_TIMESTAMP_FILE}" -o -z "${LMOD_CACHE_DIR}" ]; then
     # Get number of cache Dirs:
     local matchStr="number of cache dirs"
     numSC=$(module --config 2>&1 | grep "$matchStr" | sed -e "s/$matchStr *//")

     if [ "$numSC" = 1 ]; then
        cacheLine=$(module --config 2>&1 | grep -A 2 "Cache Directory" | tail -1 )
        dir=$(echo $cacheLine | tr -s ' ' | cut -f1 -d' ')
        ts=$( echo $cacheLine | tr -s ' ' | cut -f2 -d' ')

        if [ -z $LMOD_CACHE_DIR ]; then
           LMOD_CACHE_DIR=$dir
        fi

        if [ -z $LMOD_CACHE_TIMESTAMP_FILE ]; then
           LMOD_CACHE_TIMESTAMP_FILE=$ts
        fi
     fi
     if [ -z "${LMOD_CACHE_TIMESTAMP_FILE}" -o -z "${LMOD_CACHE_DIR}" ]; then
         warning "Failed to determine defaults via module --config"
     fi
  fi
}

# check script parameters (only in debug mode)
check_parameters() {
    # required parameters
    for param in LMOD_CACHE_DIR LMOD_CACHE_TIMESTAMP_FILE BASE_MODULEPATH
    do
        debug "\$${param}: ${!param}"
        if [ -z ${!param} ]; then
            error "\$$param is undefined"
        fi
    done
    # optional parameters
    for param in LMOD_CACHE_TIMESTAMP_FILE_TXT
    do
        debug "\$${param}: ${!param}"
    done
}

# create new cache timestamp file
new_timestamp() {
    local dir=$(dirname $LMOD_CACHE_TIMESTAMP_FILE)

    if [ ! -d $dir ]; then
       mkdir -p $dir
    fi

    debug "Cleaning up existing file ${LMOD_CACHE_TIMESTAMP_FILE}.new (if any)"
    debug "ls -l \${LMOD_CACHE_TIMESTAMP_FILE}.new: `ls -l ${LMOD_CACHE_TIMESTAMP_FILE}.new 2> /dev/null`"
    rm -f ${LMOD_CACHE_TIMESTAMP_FILE}.new

    debug "Creating Lmod cache timestamp file @ ${LMOD_CACHE_TIMESTAMP_FILE}.new"
    debug "Creating empty timestamp file at ${LMOD_CACHE_TIMESTAMP_FILE}.new"
    touch ${LMOD_CACHE_TIMESTAMP_FILE}.new
}


# install new cache file (.lua or .luac_*)
install_new_cache() {
    cache_ext=$1
    cache_file_path=$2
    chmod 644 ${cache_file_path}.new.${cache_ext}

    local OLD=${cache_file_path}.old.${cache_ext}
    local NEW=${cache_file_path}.new.${cache_ext}
    local RST=${cache_file_path}.${cache_ext}

    # back up existing cache file
    if [ -f $RST ]; then
        debug "Backup up old Lmod cache file $RST to $OLD"
        cp $RST $OLD
        debug "`ls -l $OLD`"
    fi

    # put new cache file in place
    debug "Putting new cache file $NEW place at $RST"
    mv -f $NEW $RST
}

# run spider command to generate specified cache type to specified file path
run_spider() {
    cache_type=$1
    cache_file_path=$2

    if [ -z "$LMOD_DEBUG" ]; then
      debug "Generating new $cache_type Lmod cache file at $cache_file_path"
      debug "Running \"spider -o $cache_type $BASE_MODULEPATH > $cache_file_path\""
      /usr/share/lmod/lmod/libexec/spider -o $cache_type $BASE_MODULEPATH > $cache_file_path
    else
      /usr/share/lmod/lmod/libexec/spider -D -o $cache_type $BASE_MODULEPATH > $cache_file_path 2> ${cache_file_path}.log
    fi

    SPIDER_EXIT_CODE=$?
    debug "Exit code of spider command: $SPIDER_EXIT_CODE"

    if [ -n "$PYTHON" ]; then
      if [ $cache_type = "jsonReverseMapT" -o $cache_type = "xalt_rmapT" ]; then
        mv  $cache_file_path $cache_file_path.new
        $PYTHON -mjson.tool $cache_file_path.new > $cache_file_path
        rm $cache_file_path.new
      fi
    fi

    if [ $SPIDER_EXIT_CODE -ne 0 ]; then
        error "Failed to update Lmod spider cache (exit code: $SPIDER_EXIT_CODE)"
    fi
}

# update Lmod cache files (both .lua and .luac_*)
update_cache() {
    cache_dir=$1
    cache_type=$2
    cache_file_name=$cache_dir/${cache_type}
    ext=lua
    if [ "$cache_type" = "jsonReverseMapT" -o "$cache_type" = "xalt_rmapT" ]; then
      ext=json
    fi

    if [ ! -d  $cache_dir ]; then
       mkdir -p $cache_dir
    fi

    debug "Updating $cache_type Lmod cache at $cache_file_name"

    # generate new plain-text Lmod cache file (.lua)
    run_spider $cache_type ${cache_file_name}.new.$ext
    install_new_cache $ext ${cache_file_name}

    if [ $ext = lua ]; then
       # also install compiled Lmod cache version
       # include Lua version in filename of compiled cache, to avoid compatibility issues when Lua is upgraded
       lua_ver=$(/usr/bin/lua -e 'print((_VERSION:gsub("Lua ","")))')
       luac -o ${cache_file_name}.new.luac_$lua_ver ${cache_file_name}.lua
       install_new_cache luac_$lua_ver ${cache_file_name}
    fi
}

#########################################################################
#
# MAIN
#
#########################################################################

parse_cmdline $@

# put defaults in place for parameters that were not set
set_defaults

# check parameters (make sure they are defined)
check_parameters

# See if there is a working python
for cmd in python3 python python2; do
  if command -v $cmd > /dev/null ; then
    PYTHON=$cmd
    break
  fi
done

if [ -z "$LMOD_REVERSEMAPT_DIR" ]; then
  LMOD_REVERSEMAPT_DIR="$LMOD_CACHE_DIR/../reverseMapD"
fi

# create new timestamp file
# timestamp file marks time of last change to the system, and *must* be older than the Lmod cache
# files themselves a timestamp file that is more recent than the Lmod cache file(s) indicates
# cache invalidity
new_timestamp

# update cache files
update_cache $LMOD_CACHE_DIR spiderT
if [ $UPDATE_REVERSEMAPT_CACHE -ne 0 ]; then
    update_cache  $LMOD_REVERSEMAPT_DIR reverseMapT
fi

if [ "$UPDATE_JSON_RMAPT" -ne 0 ]; then
    update_cache $LMOD_REVERSEMAPT_DIR jsonReverseMapT
fi

if [ "$UPDATE_XALT_JSON_RMAPT" -ne 0 ]; then
    update_cache $LMOD_REVERSEMAPT_DIR xalt_rmapT
fi

# put new timestamp file in place
debug "put new timestamp file in place"
mv ${LMOD_CACHE_TIMESTAMP_FILE}.new $LMOD_CACHE_TIMESTAMP_FILE
