#!/bin/bash
SCRIPT_LOCATION="$( readlink -f "${BASH_SOURCE[${#BASH_SOURCE[@]} - 1]}" )"
PROG_DIR="$( dirname "$SCRIPT_LOCATION")"
PROG_NAME="$(basename "$SCRIPT_LOCATION")"
JOBID=
# msg "message string" "message type" "true/false - for output to stderr"
################################################################################
# A wrapper around echo to respect verbosity                                   #
################################################################################
msg() {
    local string="$1"
    local type="${2:-info}"
    local stderr="${3:-false}"
    # local level="${4:-1}"

    string="[$type] $string"

    if [[ $FLAGS_verbose -eq 0 ]]; then
        case "$stderr" in
            "false")
                echo "$string"
                ;;
            "true")
                echo "$string" 1>&2
                ;;
            *)
                error $LINENO "unknown STDERR option: '$stderr'"
                ;;
        esac
    fi    
}


logjob() {
    output=$("$@")
    JOBID=$(echo "$output" | perl -ne '$_ =~ m/Your\s+job[ -]array\s+(\d+)\..+submitted$/i; print $1')
}


# log_a_file <file_to_log> <where_to_log> <where to write the log> <OPTIONAL MD5>
################################################################################
#                                           #
################################################################################
log_a_file() {
    local file_to_log="$1"
    local where_to_log="$2"
    local write_log_to="$3"
    local md5="$4"

    if [ -z $md5 ]; then
        md5=$(md5sum "$file_to_log")
    fi

    local file_name="$where_to_log"/"$RUN_NAME"
    cp "$file_to_log" "$file_name"
    echo "$md5 $file_name" > "$write_log_to"
}

internal_md5_check() {
    local log_file="$1"
    local check_name="$2"

    msg "checking: $log_file"
    local exp_hash=$(head -qn1 $log_file | awk '{print $1}')
    local hash_file=$(head -qn1 $log_file | awk '{print $2}')
    # local EXP_MD5=$(head -qn1 "$log_file" | cut -f1)
    # local FILE_PATH=$(head -qn1 "$log_file" | cut -f2)

    cur_hash=$(md5sum "$hash_file" | awk '{print $1}')

    if [ "$cur_hash" != "$exp_hash" ]; then
           error $LINENO "current $check_name md5 ($cur_hash) != expected md5 ($exp_hash)"
    else
        msg "$check_name md5 ok"
    fi

} 


# error  <LINENO> <ERROR_MSG>
################################################################################
# called by a trap when ERR or SIGINT                                          #
################################################################################
error() {
    local lineno=$1
    local error="${2:-error!}"
    
    FLAGS_verbose=0
    msg "$error" "error" "true"

    # restore files to their previous run state
    restore

    # removes run log directory
    cleanup

    # remove backup directory
    cleanup_bak

    # delete any temp files that have been created during the run
    delete_temp_exit

    # delete any temp files that have been created during the run but are 
    # unfinished
    delete_temp_error

    # The last thing we do is to delete the run state file
    delete_run_state
    exit 1
}


# exit clean
################################################################################
# called when the script exits cleanly                                         #
################################################################################
exit_clean() {
    # msg "clean exit" "exit"

    # delete any temp files that have been created during the run
    delete_temp_exit

    # remove as no longer required
    cleanup_bak

    # The last thing we do is to delete the run state file
    delete_run_state
    exit 1
}


# Function Description
################################################################################
# This performs a backup of all the files that will be changed during a run
################################################################################
base_backup() {
    for i in "${!EXP_FILES[@]}"; do
        # echo "${i}=""${EXP_FILES[$i]}"
        backup "$i" "${EXP_FILES[$i]}"
    done
}


# backup <NAME> <SOURCE_PATH>
################################################################################
# This copies a file from the source path to the backup directory $BACKUP_DIR  #
################################################################################
backup() {
    # The name forms akey in a backup/restore associative arrays
    local name="$1"
    local path_to_file=$(readlink -f "$2")
    # echo "BACKUP=$name:""$path_to_file"
    # generate a backup file name
    local file_name=$(basename "$path_to_file")
    local backup_file="$BACKUP_DIR"/"$file_name"

    # log the original path and the backup path in restore and backup 
    # associative arrays
    BACKUPS["$name"]="$path_to_file"
    RESTORE["$name"]="$backup_file"

    # now backup the file
    cp "$path_to_file" "$BACKUP_DIR"/"$file_name"
}


# backup <NAME> <SOURCE_PATH>
################################################################################
# This copies a file from the source path to the backup directory $BACKUP_DIR  #
################################################################################
restore() {
    for i in "${!RESTORE[@]}"; do
        # Make sure there is something to copy first
        if [ -z "${RESTORE[$i]}" ] || [ -z "${BACKUPS[$i]}" ]; then
            # echo "${RESTORE[$i]}"
            # echo "RESTORING $i="${BACKUP[$i]}"" TO:
            msg "empty slot: $i unable to restore" "error" "true"
        else
            msg "restoring $i to ${BACKUPS[$i]}"
            cp "${RESTORE[$i]}" "${BACKUPS[$i]}" 
        fi
    done
}


# cleanup
################################################################################
# Called on error to remove run log directory                                  #
################################################################################
cleanup() {
    if [ ! -z "$RUN_LOG" ] && [ -e "$RUN_LOG" ] && [ -d "$RUN_LOG" ]; then
        echo "[cleanup] removing "$RUN_LOG""
        rm -r "$RUN_LOG"
    fi
}


# cleanup_bak
################################################################################
# Called on error to remove run .bak directory                                 #
################################################################################
cleanup_bak() {
    if [ ! -z "$BACKUP_DIR" ] && [ -e "$BACKUP_DIR" ] && [ -d "$BACKUP_DIR" ]; then
        msg "removing backup: $BACKUP_DIR"
        rm -r "$BACKUP_DIR"
    fi
}


# delete_run_state
################################################################################
# This deletes the running state file. This should be the last thing done      #
# before clean exit or error exit                                              #
################################################################################
delete_run_state() {
    # The last thing we do is to delete the run state file
    [ -e "$RUN_STATE" ] && rm "$RUN_STATE"
}


# delete_temp_error
################################################################################
# Called when the program exits through error                                  #
################################################################################
delete_temp_error() {
    for i in "${TEMP_ERR[@]}"; do
        msg "deleting '$i'..." "cleanup"
        rm "$i"
    done
}


# delete_temp_exit
################################################################################
# Called when the program exits through error or cleanly                       #
################################################################################
delete_temp_exit() {
    for i in "${TEMP_EXIT[@]}"; do
        msg "deleting '$i'..." "cleanup"
        rm "$i"
    done
}


# This should be sourced into a job array run script after shflags has been 
# loaded and you have defined any of your command line arguments
DEFINE_string 'jafile' '' 'a file with job array parameters, input file is first column, output file is second column' 't'
DEFINE_string 'force' '' 'for the running of the given job array file with minimal checks' 'f'
DEFINE_boolean 'update' false 'if any of the parameters of the run have changed then this should be used to update the run' 'u'
DEFINE_boolean 'new' false 'if this is the first run of the job then the new flag should be given for initial setup' 'n'
DEFINE_integer 'local_job_id' 1 'the SGE_TASK_ID to mimic for a --nocluster job'
DEFINE_boolean 'verbose' false 'do you want to print some info' 'v'
DEFINE_boolean 'cluster' true 'if true then the job is run on the cluster, if --nocluster is used then as single job is run on the head node with /usr/bin/time -v, this can be useful for testing/debugging'
DEFINE_boolean 'auto' true 'if true then the qsub command is run for you with --memory, --shell, --time, --batch, --step. If you want to define your own qsub the use --noauto'
DEFINE_integer 'profile' -1 'run a small subset of jobs to try to determine memory usage' 'p'
DEFINE_integer 'batch' -1 'the max number of concurrent jobs that will run at the same time, the default is -1 (so no task concurrency is enforced)'
DEFINE_integer 'step' 1 'the step size for the job'
DEFINE_string 'mem' '0' 'the memory requirments for the job'
DEFINE_string 'shell' '/bin/bash' 'the shell to use for the job'
DEFINE_string 'time' '00:05:00' 'the max time for the job'
DEFINE_string 'nodes' '' 'The node list for the job'
DEFINE_string 'scratch' '' 'if scratch space is used then /scratch0 should be enabled and used in scripts'

# DEFINE_float 'long_opt' 'default_opt' 'description' 'short_opt'
# DEFINE_boolean 'long_opt' 'default_opt' 'description' 'short_opt'

FLAGS_HELP=${USAGE}

# parse the command-line
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

# This is the working directory where all run information will be stored
WORKING_DIR="$1"

# trap errors and CTRL-C 
trap "error $LINENO 'error...'" ERR SIGINT
trap "exit_clean" EXIT

# if the script location is not set then
if [ -z $SCRIPT_LOCATION ]; then
    error $LINENO "need to know the script location and can't find it!"
fi

# if this is a --new run then --force and --update can't be set
if [[ $FLAGS_new -eq 0 ]]; then
     if [ ! -z $FLAGS_force ] || [ $FLAGS_update -eq 0 ]; then
         error $LINENO "--force or --update can't be used with --new"
     fi
fi

# If we are trying to force a job array then we must make sure we have a 
# valid file
if [ ! -z $FLAGS_force ] && [ ! -r $FLAGS_force ]; then
    error $LINENO "the --force job array file is not readable"
fi

if [[ ! -z $FLAGS_jafile ]] && [[ ! -r "$FLAGS_jafile" ]]; then
    error $LINENO "job array file does not exist or is not readable: $FLAGS_jafile"
fi

# Some tests on the working directory
if [ -z "$WORKING_DIR" ]; then
    error $LINENO "working directory not defined"
fi

if [ ! -e "$WORKING_DIR" ] && [ $FLAGS_new -eq 1 ]; then
    error $LINENO "working directory does not exist and --new is not set"
fi

if [ -e "$WORKING_DIR" ] && [ ! -d "$WORKING_DIR" ]; then
    error $LINENO "working directory is not a directory"
fi

# Only while debugging
##rm -rf "$WORKING_DIR"

if [ -d "$WORKING_DIR" ] && [ $FLAGS_new -eq 0 ]; then
    error $LINENO "working directory exists and --new is set"
fi

# These are files that will be deleted on clean exit
TEMP_EXIT=()

# These are files that will be deleted on error exit
TEMP_ERR=()

# We check the root of the working directory to make sure we can
# write to it and create the working directory
WORKING_DIR_ROOT=$(dirname "$WORKING_DIR")
if [[ ! -e "$WORKING_DIR_ROOT" ]] || [[ ! -w "$WORKING_DIR_ROOT" ]]; then
    error $LINENO "the working direcotry root '$WORKING_DIR_ROOT' does not exist or is not writable!"
fi

# If we get here then initial checks on the working directory are ok
# so we can start to do stuff
# first get the full path to the working directory
WORKING_DIR=$(readlink -f "$WORKING_DIR")

SCRIPT_MD5=$(md5sum "$SCRIPT_LOCATION" | awk '{print $1}')

# These are files and directories that will either exist or will be
# created if it is a new run
BACKUP_DIR="$WORKING_DIR"/".bak"
RUN_STATE="$WORKING_DIR"/".error"

declare -A EXP_FILES=()
EXP_FILES[PATHSPEC]="$WORKING_DIR"/".pathspec"
EXP_FILES[RUN_NO]="$WORKING_DIR"/".run_no"
EXP_FILES[LAST_RUN]="$WORKING_DIR"/".last_run"
EXP_FILES[RUN_LOG]="$WORKING_DIR"/"run.log"
EXP_FILES[CUR_JA]="$WORKING_DIR"/".current_job_array"
EXP_FILES[CUR_SCRIPT]="$WORKING_DIR"/".current_script"

declare -A EXP_DIR=()
EXP_DIR[JOB_ARRAY_DIR]="$WORKING_DIR"/".jobarray"
EXP_DIR[JOB_FORCED_DIR]="$WORKING_DIR"/".jobforced"
EXP_DIR[SCRIPTS_DIR]="$WORKING_DIR"/".scripts"

RUNTIME=$(date +%s)
DATE_PRETTY=$(date)
PROCESS=$$

declare -A BACKUPS=()
declare -A RESTORE=()
#declare -A DELETES=()

msg "*** $PROG_NAME v$VERSION ***"
msg "working directory: $WORKING_DIR"
msg "script location: $SCRIPT_LOCATION"
msg "script md5: $SCRIPT_MD5"
msg "process ID: $RUNTIME"
msg "runtime: $RUNTIME"

# First we check for an error state file
if [ -e "$RUN_STATE" ]; then
    error "$LINENO" "error state found: this means that we didn't exit cleanly. manually check everything, if ok then delete: $RUN_STATE"
fi

ALERT="-"
# if this is a new run then we need to do some initialisation
if [[ $FLAGS_new -eq 0 ]]; then
    msg "initialising new run...."

    # This is the first run
    RUN_NO=0
    RUN_NAME="${RUN_NO}_${PROCESS}_${RUNTIME}"
    ALERT="I"

    # if this is a new run then a job array file needs to be set
    # if it is not then this is an error
    if [[ -z "$FLAGS_jafile" ]]; then
        error $LINENO "job array must be defined for a new job"
    fi

    FLAGS_jafile=$(readlink -f "$FLAGS_jafile")
    JAFILE_MD5=$(md5sum "$FLAGS_jafile" | awk '{print $1}')
    msg "job array file location: $FLAGS_jafile"
    msg "job array file md5: $JAFILE_MD5"

    # create dir
    mkdir "$WORKING_DIR"
    touch "$RUN_STATE"

    # Now we create the directories in the workspace
    for i in "${!EXP_DIR[@]}"; do
        msg "creating $i: ${EXP_DIR[$i]}"
        mkdir "${EXP_DIR[$i]}"
    done

    for i in "${!EXP_FILES[@]}"; do
        msg "creating $i: ${EXP_FILES[$i]}"
        touch "${EXP_FILES[$i]}"
    done
    
    # We initialise the run_no to 0
    echo "$RUN_NO" > "${EXP_FILES[RUN_NO]}"

    # We initialise the run log
    echo -e "flag\tdate\truntime\tprocess_id\trun_no\trun_name\tworking_dir\tjob_array_file\tjob_array_file_md5\tscript_location\tscript_md5\tout_err_log\t#jobs\t#run\t#test_fails\tjobid" > "${EXP_FILES[RUN_LOG]}"
    echo -e "$ALERT\t$DATE_PRETTY\t$RUNTIME\t$PROCESS\t$RUN_NO\tINIT\t$WORKING_DIR\t$FLAGS_jafile\t$JAFILE_MD5\t$SCRIPT_LOCATION\t$SCRIPT_MD5\t\t\t\t\t" >> "${EXP_FILES[RUN_LOG]}"

    # Initialise the pathspec, this will be checked to make sure we are running
    # from the same location
    echo "$WORKING_DIR" > "${EXP_FILES[PATHSPEC]}"

    # Now we take a copy of the jobarray and the script and store them
    # TODO: these probably should be deleted on error
    log_a_file "$FLAGS_jafile" "${EXP_DIR[JOB_ARRAY_DIR]}" "${EXP_FILES[CUR_JA]}" "$JAFILE_MD5" 
    log_a_file "$SCRIPT_LOCATION" "${EXP_DIR[SCRIPTS_DIR]}" "${EXP_FILES[CUR_SCRIPT]}" "$SCRIPT_MD5"

    # Now that the job array is logged we zero the variable so it does not
    # interfere with the --jafile --update options below
    FLAGS_jafile=
    ALERT="-"
else
    touch "$RUN_STATE"
fi

# Now we do some checks of the
# pathspec ok
PATHSPEC=$(head -qn1 "${EXP_FILES[PATHSPEC]}")

if [ "$PATHSPEC" != "$WORKING_DIR" ]; then
    error $LINENO "pathspec: $PATHSPEC != working directory $WORKING_DIR...have you moved the working directory?"
else
    msg "pathspec ok"
fi

# ls -lah "$WORKING_DIR"

# internal jobarray ok
internal_md5_check "${EXP_FILES[CUR_JA]}" "job array"

# script ok
#internal_md5_check "${EXP_FILES[CUR_SCRIPT]}" "script"

# Create the backup direcotry
mkdir "$BACKUP_DIR"

# create backups of files to be changed
base_backup

# Update files
# first get the run number and increment
RUN_NO=$(head -qn1 "${EXP_FILES[RUN_NO]}")
RUN_NO=$((RUN_NO + 1))
echo "$RUN_NO" > "${EXP_FILES[RUN_NO]}"

RUN_NAME="${RUN_NO}_${PROCESS}_${RUNTIME}"
msg "current run number: $RUN_NO"
msg "current run name: $RUN_NAME"


if [[ ! -z $FLAGS_jafile ]] && [[ $FLAGS_update -eq 1 ]]; then
    # changing the job array file without specifying --force - error
    error $LINENO "if you want to update a job array file use --update in addition to --jafile"
elif [[ ! -z $FLAGS_jafile ]] && [[ $FLAGS_update -eq 0 ]]; then
    # Changing the job array file
    msg "updating the job array file to: $FLAGS_jafile"
    ALERT="J"
    FLAGS_jafile=$(readlink -f "$FLAGS_jafile")
    JAFILE_MD5=$(md5sum "$FLAGS_jafile" | awk '{print $1}')
    msg "new job array file location: $FLAGS_jafile"
    msg "new job array file md5: $JAFILE_MD5"
    log_a_file "$FLAGS_jafile" "${EXP_DIR[JOB_ARRAY_DIR]}" "${EXP_FILES[CUR_JA]}" "$JAFILE_MD5"
    # cat "${EXP_FILES[CUR_JA]}"
    # We want to delete the logged copy of the job array file should we
    # error out as we want to return to the original state
    logged_ja_file=$(head -qn1 "${EXP_FILES[CUR_JA]}" | awk '{print $2}')
    msg "adding $logged_ja_file to temp err"
    TEMP_ERR+=("$logged_ja_file")
fi

# has the user changed ja without force - error
# has the user changed ja with force backup old assign new (workout a way of restotring
CUR_SCRIPT=$(head -qn1 "${EXP_FILES[CUR_SCRIPT]}" | awk '{print $2}')
CUR_SCRIPT_MD5=$(head -qn1 "${EXP_FILES[CUR_SCRIPT]}" | awk '{print $1}')

if [ "$SCRIPT_MD5" != "$CUR_SCRIPT_MD5" ]; then
    if [[ $FLAGS_update -eq 1 ]]; then
        # If we are not updating this is an error, it could be the user has
        # run the wrong script
        error $LINENO "The md5 of the script that is running is not the same as the logged script, if this is intentional run with --update"
    else
        ALERT="S"
        # the user has intentionally updated the script
        msg "new script used: $SCRIPT_MD5"
        log_a_file "$SCRIPT_LOCATION" "${EXP_DIR[SCRIPTS_DIR]}" "${EXP_FILES[CUR_SCRIPT]}" "$SCRIPT_MD5" 

        logged_script_file=$(head -qn1 "${EXP_FILES[CUR_SCRIPT]}" | awk '{print $2}')
        #CUR_SCRIPT_MD5=$(head -qn1 "${EXP_FILES[CUR_SCRIPT]}" | awk '{print $1}')
        msg "adding $logged_script_file to temp err"
        TEMP_ERR+=("$logged_script_file")
    fi
else
    msg "script md5 matches"
fi

# # Update files
# # first get the run number and increment
# RUN_NO=$(head -qn1 "${EXP_FILES[RUN_NO]}")
# RUN_NO=$((RUN_NO + 1))
# store the nun number
# echo "$RUN_NO" > "${EXP_FILES[RUN_NO]}"

# RUN_NAME="${RUN_NO}_${PROCESS}_${RUNTIME}"

# msg "current run number: $RUN_NO"
# msg "current run name: $RUN_NAME"

RUN_LOG="$WORKING_DIR"/"$RUN_NAME"
RUN_ERR_LOG="$RUN_LOG"/"err"
RUN_OUT_LOG="$RUN_LOG"/"out"
RUN_FILE_LOG="$RUN_LOG"/"file"
RUN_FILE_PROFILE="$RUN_LOG"/"profile" # created on demand

# These are the locations of the job STDERR and STDOUT files
# they should be used in your qsub arguments -e and -o
ERRORS="$RUN_ERR_LOG"/\$TASK_ID_"${RUN_NAME}".err
OUTS="$RUN_OUT_LOG"/\$TASK_ID_"${RUN_NAME}".out

#dirs_to_create=
for i in "$RUN_LOG" "$RUN_ERR_LOG" "$RUN_OUT_LOG" "$RUN_FILE_LOG"; do
    msg "creating log dir: $i"
    mkdir "$i"    
done

INIT_RUN_FILE="$RUN_FILE_LOG"/"${RUN_NAME}"_initial
FINAL_RUN_FILE="$RUN_FILE_LOG"/"${RUN_NAME}"_final


FORCE_RUN=1
#echo "$FLAGS_jafile"
FLAGS_jafile=$(head -qn1 "${EXP_FILES[CUR_JA]}" | awk '{print $2}')
#echo "$FLAGS_jafile"
# Do I need to implement no checks?
if [ ! -z $FLAGS_force ]; then
    FORCE_RUN=0
    ALERT="F"
    msg "forcing the job array: $FLAGS_force"
    FLAGS_jafile=$FLAGS_force
    cp "$FLAGS_force" "$INIT_RUN_FILE"
else
    #logged_ja_file=$(head -qn1 "${EXP_FILES[CUR_JA]}" | awk '{print $2}')
    msg "copy $FLAGS_jafile to $INIT_RUN_FILE"
    cp "$FLAGS_jafile" "$INIT_RUN_FILE"
fi 

JAFILE_MD5=$(head -qn1 "${EXP_FILES[CUR_JA]}" | awk '{print $1}')

# finally we initialise the verbose variable so it can be used as an argument
# to other scripts in order that they can respect the --verbose argument to
# the run script
VERBOSE=
if [[ $FLAGS_verbose -eq 0 ]]; then
    VERBOSE="-v"
fi

TEMP="/tmp"
if [ ! -z $FLAGS_scratch ] && [[ $FLAGS_cluster -eq 0 ]]; then
    TEMP="/scratch0"
fi

msg "initialisation complete"
 
# Only while debugging
#rm -r "$WORKING_DIR"
