#!/bin/bash

# Copyright (C) 2017 Scibian Project <legal@scibian.org>
#
# This file is part of scibian-singularity-slurm.
#
# scibian-singularity-slurm is free software: you can redistribute it
# and/or modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# scibian-singularity-slurm is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied warranty
# of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with scibian-singularity-slurm.  If not, see
# <http://www.gnu.org/licenses/>.

set -ueo pipefail

srun_legacy=${SINGULARITY_SRUN_LEGACY:-/usr/bin/srun.legacy}

# If not into a container then exec legacy ssh
if [ -z "${SINGULARITY_CONTAINER:-}" ]; then
  exec "$srun_legacy" "$@"
fi

slurm_tmpdir="${SINGULARITY_SLURM_TMPDIR:-$HOME/.singularity_slurm_tmp}"
mkdir -p "$slurm_tmpdir"

readonly SRUN_SINGLE_OPTS="[EHIkKlOQsuvXZhV]"
readonly SRUN_MULTI_OPTS="[AbcdDeiJLMmnNopqrStTWCBG]"
readonly SRUN_MULTI_LONG_OPTS="
account acctg-freq extra-node-info bb begin bxi-bandwidth cluster-constraint comment
constraint cores-per-socket cpu-bind cpu-freq cpus-per-gpu cpus-per-task
deadline delay-boot dependency chdir error epilog export nodefile gid gpus
gpu-bind gpu-freq gpus-per-node gpus-per-socket gpus-per-task gres hint
input job-name jobid licenses clusters distribution mail-type mail-user
msc-label mem mem-per-cpu mem-per-gpu mem-bind mincpus msg-timeout mpi
nodes ntasks network ntasks-per-core ntasks-per-node ntasks-per-socket
output open-mode het-group partition power priority profile prolog qos
relative reservation core-spec signal slurm-debug sockets-per-node switches
threads time task-epilog task-prolog thread-spec threads-per-core time-min
tmp uid wait nodelist wckey exclude x11
"

SRUN_OPTIONS=()
SRUN_EXE=

# While an option is found, stack it
# Stack script parameter as well, and the options just after it
# Stop at the first parameter which is not an option nor the script parameter
while (( $# > 0 )) ;do
  case "${1:-}" in
    --*)
      if [[ "$(echo $SRUN_MULTI_LONG_OPTS)" =~ (^| )"${1#--}"( |$) ]]; then
        SRUN_OPTIONS+=("$1" "${2:-}")
        shift 2
      else
        SRUN_OPTIONS+=("$1")
        shift
      fi
      ;;
    -*)
      if [[ "$1" =~ -${SRUN_SINGLE_OPTS}*$ ]]; then
        SRUN_OPTIONS+=("$1")
        if [[ "$1" =~ ^.*I$ ]]; then
          if [ "${2:-}" =~ ^[0-9]+$ ]]; then
            SRUN_OPTIONS+=("$2")
            shift
          fi
        fi
	shift
      elif [[ "$1" =~ -${SRUN_SINGLE_OPTS}*${SRUN_MULTI_OPTS}$ ]]; then
        SRUN_OPTIONS+=("$1" "${2:-}")
        shift 2
      else
        # Do not stack invalid options (will be handled by legacy slurm command)
        break
      fi
      ;;
    *)
      SRUN_EXE="${1:-}"
      shift
      break
      ;;
  esac
done

if [ -z "${SRUN_EXE:-}" ]; then
  # No script -> call legacy srun
  if (( ${#SRUN_OPTIONS[@]} > 0 )); then
    exec "$srun_legacy" "${SRUN_OPTIONS[@]}" "$@"
  else
    exec "$srun_legacy" "$@"
  fi
fi

if [ "${SRUN_EXE:0:1}" != "." ] && [[ "${SRUN_EXE}" =~ ^[^/]+$ ]]; then
  if found="$(which "$SRUN_EXE")"; then
    SRUN_EXE="$found"
  fi
fi
if [ -f "$SRUN_EXE" ]; then
  SRUN_EXE="$(realpath "$SRUN_EXE")"
fi

tmpscript=$(mktemp --tmpdir="$slurm_tmpdir" "$(basename "$SRUN_EXE").XXX")
trap 'rm -f $tmpscript' EXIT

echo "#!/bin/sh" >$tmpscript
if file -b "$SRUN_EXE" | grep -qw text; then
  sed -E -e '1d' -e '/^(#.*|\s*)$/!Q' "$SRUN_EXE" >>$tmpscript
  # Retrieve interpretor for when the script is not executable
  if ! [ -x "$SRUN_EXE" ]; then
    interp=$(head -1 "$SRUN_EXE" | grep '^#!') || interp=/bin/sh
    interp="${interp#\#!}"
  fi
fi
echo "export SINGULARITY_OPTIONS='${SINGULARITY_OPTIONS:-}'" >>$tmpscript
# Option -s: see https://github.com/hpcng/singularity/issues/3250#issuecomment-481323679
echo "exec singularity -s exec ${SINGULARITY_OPTIONS:-} '$SINGULARITY_CONTAINER' ${interp:-} '$SRUN_EXE' \"\$@\"" >>$tmpscript
chmod a+x $tmpscript

# PATH and LD_LIBRARY_PATH should be kept
export SINGULARITYENV_PATH="$PATH"
export SINGULARITYENV_LD_LIBRARY_PATH="$LD_LIBRARY_PATH"

if (( ${#SRUN_OPTIONS[@]} > 0 )); then
  "$srun_legacy" "${SRUN_OPTIONS[@]}" $tmpscript "$@"
else
  "$srun_legacy" $tmpscript "$@"
fi
