#!/usr/bin/env bash
set -euo pipefail

function usage() { cat << EOF
cacher: RENEE job submission wrapper script for caching remote file or resources.
USAGE:
  cacher <MODE> [OPTIONS] -s SIF_CACHE -i IMAGE_URIS -t TMP_DIR
SYNOPSIS:
  This script submits the cacher master job to the cluster. This main process dictates
how subsequent resources are pulled onto the cluster's filesystem. cacher utilizes SLURM
to avoid pull into resources on a compute node but support for additional job schedulers
(i.e. PBS, SGE, LSF, Tibanna) may be added in the near future.
  The main entry point of the pipeline 'renee' calls this job submission wrapper script.
As so, this script can be used to manually by-pass 'renee' for a previously failed cache.
  Please Note: it is highly recommended to use 'renee'; it is the main entry point
and preferred entry point of the RENEE pipeline. If you are experience error, it
maybe due to improperly mounting singularity bind paths which 'renee' will internally
handle.
Required Positional Argument:
  [1] MODE  [Type: Str] Define the snakemake executor mode.
                        Valid mode options include: <slurm>
                         a) slurm: uses slurm and singularity snakemake backend.
                             The slurm EXECUTOR will submit jobs to the cluster.
                             It is recommended running RENEE in this mode as
                             most of the steps are computationally intensive.
Required Arguments:
  -s, --sif-cache   [Type: Path]  Path to output directory to cache remote resources.
  -i, --image-uris  [Type: Str]   Image(s) to pull from Dockerhub. Multiple images
                                   are separated by a comma.
OPTIONS:
  -t, --tmp-dir     [Type: Path]  Path to tmp singularity dir. Singularity uses this
                                   directory when images are pulled from DockerHub
                                   and converted into SIFs. If not provided, the
                                   location to the temp dir will default to the
                                   following "/tmp/$USER/cacher/.singularity/"
                                   directory.
  -h, --help        [Type: Bool]  Displays usage and help information.
Example:
  $ cacher slurm \\
      -s /scratch/$USER/cache \\
      -i 'docker://nciccbr/ccbr_arriba_2.0.0:v0.0.1,docker://nciccbr/ccbr_rna:v0.0.1'
Version:
  0.1.0
EOF
}


# Functions
function err() { cat <<< "$@" 1>&2; }
function fatal() { cat <<< "$@" 1>&2; usage; exit 1; }
function abspath() { readlink -e "$1"; }
function parser() {
  # Adds parsed command-line args to GLOBAL $Arguments associative array
  # + KEYS = short_cli_flag ("j", "o", ...)
  # + VALUES = parsed_user_value ("MasterJobName" "/scratch/hg38", ...)
  # @INPUT "$@" = user command-line arguments
  # @CALLS check() to see if the user provided all the required arguments

  while [[ $# -gt 0 ]]; do
    key="$1"
    case $key in
      -h  | --help) usage && exit 0;;
      -s  | --sif-cache)  provided "$key" "${2:-}"; Arguments["s"]="$2"; shift; shift;;
      -i  | --image-uris) provided "$key" "${2:-}"; Arguments["i"]="$2"; shift; shift;;
      -t  | --tmp-dir)    provided "$key" "${2:-}"; Arguments["t"]="$2"; shift; shift;;
      -*  | --*) err "Error: Failed to parse unsupported argument: '${key}'."; usage && exit 1;;
      *) err "Error: Failed to parse unrecognized argument: '${key}'. Do any of your inputs have spaces?"; usage && exit 1;;
    esac
  done

  # Check for required args
  check
}


function provided() {
  # Checks to see if the argument's value exists
  # @INPUT $1 = name of user provided argument
  # @INPUT $2 = value of user provided argument
  # @CALLS fatal() if value is empty string or NULL

  if [[ -z "${2:-}" ]]; then
     fatal "Fatal: Failed to provide value to '${1}'!";
  fi
}


function check(){
  # Checks to see if user provided required arguments
  # @INPUTS $Arguments = Global Associative Array
  # @CALLS fatal() if user did NOT provide all the $required args

  # List of required arguments
  local required=("s" "i")
  #echo -e "Provided Required Inputs"
  for arg in "${required[@]}"; do
    value=${Arguments[${arg}]:-}
    if [[ -z "${value}" ]]; then
      fatal "Failed to provide all required args.. missing ${arg}"
    fi
  done
}


function retry() {
  # Tries to run a cmd 5 times before failing
  # If a command is successful, it will break out of attempt loop
  # Failed attempts are padding with the following exponential
  # back-off strategy {4, 16, 64, 256, 1024} in seconds
  # @INPUTS "$@"" = cmd to run
  # @CALLS fatal() if command cannot be run in 5 attempts
  local n=1
  local max=5
  local attempt=true # flag for while loop
  while $attempt; do
    # Attempt command and break if successful
    "$@" && attempt=false || {
      # Try again up to 5 times
      if [[ $n -le $max ]]; then
        err "Command failed: $@"
        delay=$(( 4**$n ))
        err "Attempt: ${n}/${max}. Trying again in ${delay} seconds"
        sleep $delay;
        ((n++))
      else
        fatal "Fatal: the command has failed after max attempts"
      fi
    }
  done
}


function _pull(){
  # Caches a remote image from DockerHub
  # INPUT $1 = Snakemake Mode of execution
  # INPUT $2 = Cache output directory
  # INPUT $3 = Singularity temp directory
  # INPUT $4 = Images to pull from DockerHub

  # Check if singularity in $PATH
  # If not, try to module load singularity as a last resort
  command -V singularity &> /dev/null || {
    command -V module &> /dev/null &&
    module purge && module load singularity
  } || fatal "Fail to find or load 'singularity', not installed on target system."

  # cache executor
  executor=${1}

  # Goto Pipeline Output directory
  # Create a local singularity cache in output directory
  # cache can be re-used instead of re-pulling from DockerHub every time
  cd "$2" && export SINGULARITY_CACHEDIR="${3}"

  # unsetting XDG_RUNTIME_DIR to avoid some unsighly but harmless warnings
  unset XDG_RUNTIME_DIR

  # Run the workflow with specified executor
  case "$executor" in
    slurm)
          # Create directory for logfiles
          for image in ${4//,/$'\t'}; do
            # Try to pull image from URI with 5 max attempt
            echo "Singularity pulling ${image}"
            retry singularity pull -F ${image}
          done
        ;;
      *)  echo "${executor} is not available." && \
          fatal "Failed to provide valid execution backend: ${executor}. Please use slurm."
        ;;
    esac
}


function main(){
  # Parses args and pulls remote resources
  # @INPUT "$@" = command-line arguments
  # @CALLS pull()

  if [ $# -eq 0 ]; then usage; exit 1; fi

  # Associative array to store parsed args
  declare -Ag Arguments

  # Positional Argument for Executor
  case $1 in
    slurm) Arguments["e"]="$1";;
    -h    | --help | help) usage && exit 0;;
    -*    | --*) err "Error: Failed to provide required positional argument: <slurm>."; usage && exit 1;;
    *) err "Error: Failed to provide valid positional argument. '${1}' is not supported. Valid option(s) are slurm"; usage && exit 1;;
  esac

  # Parses remaining user provided command-line arguments
  parser "${@:2}" # Remove first item of list
  mkdir -p "${Arguments[s]}"
  cache=$(abspath "${Arguments[s]}")
  dockers="${Arguments[i]}"

  # Setting defaults for non-required arguments
  tmp="${Arguments[t]:-/tmp/$USER/cacher/.singularity/}"
  # Expand single single quoted input
  # for singularity cache directory
  eval tmp=$(echo "$tmp")
  mkdir -p "$tmp"
  tmp=$(abspath "${tmp}")

  # Print cli options prior to running
  echo -e "RENEE cache \t$(date)"
  echo -e "Running with the following parameters:"
  for key in "${!Arguments[@]}"; do echo -e "\t${key}\t${Arguments["$key"]}"; done

  # Pull remote resources into RENEE cache
  # Cache remote image from DockerHub
  # INPUT $1 = Snakemake Mode of execution
  # INPUT $2 = Cache output directory
  # INPUT $3 = Singularity temp directory
  # INPUT $4 = Images to pull from DockerHub
  _pull "${Arguments[e]}" "$cache" "$tmp" "$dockers"

}


# Main: check usage, parse args, and run pipeline
main "$@"
