#############################################################
# A template task script for running an a single step task  #
#############################################################
# Put in a random sleep so not all jobs start at the same time, this will sleep
# between 0-30s
sleep $(( $RANDOM % 30 ))

# Careful when using pipefail with zcat, if zcat is stopped early it
# exists with an error exit code and this will kill the script.
# i.e. header=$(zcat file.txt | head -n1)
# ** I think ** you can do this
# set +o pipefail
# header=$(zcat file.txt | head -n1)
# set -o pipefail
# to get around it
# A must read
# https://vaneyckt.io/posts/safer_bash_scripts_with_set_euxo_pipefail/
# Also, don't forget set -x if you want to debug
set -Euo pipefail

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
error_msg() {
    # Output an error message
    local msg="[error] ""$1"

    # We output to STDERR and STDOUT
    echo "$msg"
    echo "$msg" 1>&2
}

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
info_msg() {
    # Output an error message
    local msg="[info] ""$1"

    # We output to STDOUT
    echo "$msg"
}

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
calc_run_time() {
    local start_seconds="$1"
    if [[ -z $start_seconds ]]; then
	info_msg "runtime: unable to calculate runtime"
    else
	local end=$(date)
	local end_seconds=$(date '+%s')

	# To get the floating point numbers easier than bc
	runtime_hours=$(awk -vstart="$start_seconds" -vend="$end_seconds" 'BEGIN{time=(end-start)/3600; printf "%.2fh", time}')
	info_msg "ended at $end"
	info_msg "runtime: $runtime_hours"
    fi
}

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
error_exit() {
    # Called if we exit with an error, takes the line number of where the error
    # took place
    local LN="$1"

    # Make sure we clean up all the temp files before exiting with an error
    # exit code
    rmtemp_files

    # Errors are output to STDERR and STDOUT
    echo "[fatal] error on or near, '$LN'" 1>&2
    echo "!!!! ERROR EXIT !!!!" 1>&2
    echo "[fatal] error on or near, '$LN'"
    echo "!!!! ERROR EXIT !!!!"

    exit 1
}

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
clean_exit() {
    # Only run the exit cleanup if we have exited cleanly, this condition has
    # been added as the EXIT trap is fireing when it is erroring out
    if [[ $? -eq 0 ]]; then
	# Move any files that need moving
	move_files

	# Called if we exit cleanly
	rmtemp_files

	# calculate the runtime
	calc_run_time "$START_SECONDS"
	echo "*** END ***"
    fi

    exit $?
}

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
move_files() {
    # Move the required temp files to their final location
    local iter=0
    if [[ "${#MOVE_FROM[@]}" -gt 0 ]]; then
        for i in "${MOVE_FROM[@]}"; do
	          if [[ -e "$i" ]]; then
	              info_msg "moving $i to ${MOVE_TO[$iter]}"
	              move_single_file "$i" "${MOVE_TO[$iter]}"
	          else
	              info_msg "unable to move: could not find '${i}'"
	          fi
            iter=$((iter + 1))
        done
    fi
}

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
move_single_file() {
    # Move the required temp files to their final location
    local from="$1"
    local to="$2"

    # Take the checksum of from
    from_md5=$(md5sum "$from" | awk '{print $1}')

    # Now move the file to i's final location
    mv "$from" "$to"

    # Take the md5 of to
    to_md5=$(md5sum "$to" | awk '{print $1}')

    if [[ "$from_md5" != "$to_md5" ]]; then
	error_msg "after moving: ${from_md5} != ${to_md5}"
	rm "$to"
    else
	info_msg "md5 ${from}: ${from_md5}"
	info_msg "md5 ${to}: ${to_md5}"
    fi
}


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
rmtemp_files() {
    # Ensure all the temp directories and files created in this script are
    # removed when it exists
    # First try to delete the working directory
    if [[ -d $WORKING_DIR ]]; then
	      info_msg "deleting $WORKING_DIR"
        rm -r "$WORKING_DIR"
    fi

    # Not strictly necessary but just in case the above fails
    for i in ${TEMPFILES[@]}; do
	      if [[ -e "$i" ]]; then
	          info_msg "deleting $i"
	          rm "$i"
	      fi
    done
}


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
check_all_infiles() {
    # Expand input file paths and make sure they exist, this acts on a global
    # INFILES array
    for i in "${!INFILES[@]}"; do
        INFILES[$i]="$(readlink -f "${INFILES[$i]}")"

        # Make sure the input file is readable
        if [[ ! -r "${INFILES[$i]}" ]]; then
	          error_msg "'${INFILES[$i]}' is not readable"
	          error_exit "$LINENO"
        fi
        info_msg "infile $i=""${INFILES[$i]}"
    done
}


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
check_all_outfiles() {
    for i in "${!OUTFILES[@]}"; do
        # If the output file already exists then error out
        if [[ -e "${OUTFILES[$i]}" ]]; then
            error_msg "[error] '${OUTFILES[$i]}' already exists" 1>&2
            error_exit "$LINENO"
        fi
    done
}


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
# If the script works, the function will be called and all the tmp files removed. 
trap 'clean_exit' EXIT

# If any command fail, or cntrl+C, then the exit function will be called. 
trap 'error_exit $LINENO' SIGINT ERR SIGTERM

# Initialise the working directory
WORKING_DIR=""

# Store all the paths to the intermediate files, these will all be deleted when
# exitting cleanly or with an error. The "" is to stop unbound errors
TEMPFILES=("")

# Files that should be moved should go into these two arrays. This will then be
# moved to their final location on exit.
MOVE_FROM=()
MOVE_TO=()

START=$(date)
START_SECONDS=$(date '+%s')

# The program name
echo "=== $0 ==="
info_msg "started at $START"
info_msg "hostname=""$HOSTNAME"

# ARGUMENT 1: job array file name
# ARGUMENT 2: the tmp location, when running a live job this will be /Scratch
#             or ~/scratch0 depending on the cluster
# ARGUMENT 3: The step size for the job array, this is supplied so the task
#             script can use it in loops in combination with SGE_TASK_ID
# ARGUMENT 4: the task index value to run. If debugging this can be provided
#             at the command line. When running a live tasl then it will be the
#             SGE_TASK_ID environment variable of the currently running task

JA="$1"
TMP_LOC="$2"
STEP=$3
IDX=${4:-$SGE_TASK_ID}

info_msg "job array file=$JA"
info_msg "temp location=""$TMP_LOC"
info_msg "step=""$STEP"

# Make sure the temp directory is represented as a full path
JA="$(readlink -f "$JA")"
TMP_LOC="$(realpath "$TMP_LOC")"

info_msg "job array file=$JA"
info_msg "temp location=""$TMP_LOC"


# Make sure the SGE_TASK_ID is set, this should not be necessary with set -u
# enabled but just in case it is turned off for any reason
if [[ -z $IDX ]]; then
    error_msg "SGE_TASK_ID is null"
    error_exit "$LINENO"
fi

# Make sure the job array file is set
# If $JA does not exist or it is not readable
if [[ -z "$JA" ]] || [[ ! -r "$JA" ]]; then
    error_msg "'$JA' job array file not given or is not readable"
    error_exit "$LINENO"
fi

info_msg "job idx line (SGE_TASK_ID)="$IDX

# A template for the temp working directory that will be created to hold the
# output
TEMPLATE=".RUN_${USER}_${IDX}_XXXXXXXXXX"
info_msg "working directory template="$TEMPLATE

WORKING_DIR=$(mktemp -d -p"${TMP_LOC}" "$TEMPLATE")
info_msg "working directory=""$WORKING_DIR"

STOP=$(( IDX + STEP ))

while [[ $IDX -lt $STOP ]]; do
	  PULL_LINE=$((IDX + 1))
	  QUIT_LINE=$((PULL_LINE + 1))
	  info_msg "pull line (accounting for header)="$PULL_LINE
	  info_msg "quit line="$QUIT_LINE
	  PROCESS_LINE=($(sed -n "${PULL_LINE}p;${QUIT_LINE}q" "$JA"))
	  info_msg "process line length="${#PROCESS_LINE[@]}
	  if [[ ${#PROCESS_LINE[@]} -eq 0 ]]; then
		    break
	  fi
    ROW_IDX="${PROCESS_LINE[0]}"
    # Make sure that the the IDX == the ROWIDX it should do in all cases
    if [[ $IDX -ne $ROW_IDX ]]; then
        error_msg "idx (SGE_TASK_ID) [${IDX}] != rowidx [${ROW_IDX}]"
        error_exit "$LINENO"
    fi

    info_msg "idx (SGE_TASK_ID+STEP) [${IDX}] == rowidx [${ROW_IDX}]"

    run_command
	  IDX=$((IDX + 1))
done
