#!/usr/bin/env bash
#
# orthologer - Pipeline for deducing orthology from fasta files.
# Copyright (C) 2005-2023 E. Zdobnov lab: F. Tegenfeldt
# <fredrik.tegenfeldt@unige.ch> based on code by E. Kriventseva
#
# This file is part of orthologer.
#
# orthologer 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.
#
# orthologer 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 orthologer.  If not, see <http://www.gnu.org/licenses/>.
#
#-----------------------------------------------------
# general setup independent of batch job engine
#-----------------------------------------------------

#--- check if environment is ok -----------------
[[ "${LOADED_COMMON_PIPELINE_SH+1}" == "" ]] && echo "ERROR $0 : main config has not been loaded" 1>&2 && exit 255
#------------------------------------------------

#-----------------------------------------------------
# set global variables defined in scheduler_common.sh
#-----------------------------------------------------
CMD_SBATCH=$(/usr/bin/which sbatch 2>/dev/null)
CMD_SRUN=$(/usr/bin/which srun 2>/dev/null)
CMD_SQUEUE=$(/usr/bin/which squeue 2>/dev/null)
CMD_SCANCEL=$(/usr/bin/which scancel 2>/dev/null)
SJ_BATCH_JOBID_NAME="SLURM_JOB_ID"

sj_scheduler_init() {
    [[ ${#CMD_SQUEUE} -gt 0 ]] || { error_msg "Slurm scheduler -> missing 'squeue' command"; return 255; }
    [[ ${#CMD_SBATCH} -gt 0 ]] || { error_msg "Slurm scheduler -> missing 'sbatch' command"; return 255; }
    [[ ${#CMD_SRUN}   -gt 0 ]] || { error_msg "Slurm scheduler -> missing 'srun' command";   return 255; }
    [[ ${#CMD_SCANCEL} -gt 0 ]] || { error_msg "Slurm scheduler -> missing 'scancel' command";   return 255; }
    return 0
}

#-----------------------------------------------------
# submit the job - required by sj_submit()
#-----------------------------------------------------
sj_submit_job() {
    # create full names of stdout and stderr
    local stdout="${SJ_JOBSTDOUT}_%j.log"
    local stderr="${SJ_JOBSTDERR}_%j.log"

    # do the submission
    local jobstr=
    local ec=0
    local ntries=10 # try 10 times
    local tsleep=5  # sleep 5 secs between trials
    local submitted=0
    while (( ntries > 0 )); do
        jobstr=$($CMD_SBATCH --cpus-per-task=$SJ_NCPUPROC --wait-all-nodes=1 --no-requeue -J "$SJ_JOBLABEL" --output="$stdout" --error="$stderr" $SJ_SCHEDOPTS < "$SJ_BATCH_SCRIPT")
        ec=$?
	    #---------------------------------------------------------------------------------------
	    # make the submission robust wrt sbatch errors like:
	    # sbatch: error: Batch job submission failed: Socket timed out on send/recv operation
	    #---------------------------------------------------------------------------------------
	    # check error code - if non-zero, retry after a should slumber
	    [[ $ec -eq 0 ]] || { warn_msg "failed submitting job - retrying in $tsleep seconds"; ntries=$((ntries-1)); sleep $tsleep; continue; }
        
	    # job submitted OK - clean up and make sure loop is terminated
	    ntries=0
	    submitted=1
	    rm -f "$SJ_BATCH_SCRIPT"

	    #if [[ $ec -ne 0 ]]; then
        #    error_msg "$SJ_JOBLABEL : slurm job submission failed with error code $ec, sbatch output: $jobstr"
        #    return $ec
        #fi
	
    done
    if (( submitted == 0 )); then
        error_msg "$SJ_JOBLABEL : slurm job submission failed with error code $ec, sbatch output: $jobstr"
        return $ec
    fi
    # parse jobid
    SJ_BATCH_JOBID=$(echo $jobstr | grep "Submitted batch job" | awk '{ print $4; }')
    ec=$?
    if [[ $ec -ne 0 ]]; then
        error_msg "unable to parse job submit string from slurm (ec=$ec)"
        echo_err "---"
        echo_err "$jobstr"
        echo_err "---"
        fatal_msg "aborting!"
    fi
}

#-----------------------------------------------------
# run interactive job
#-----------------------------------------------------
sj_interactive_job() {
    # create full names of stdout and stderr
    # %j expands to job id
    local stdout="${SJ_JOBSTDOUT}_%j.log"
    local stderr="${SJ_JOBSTDERR}_%j.log"
    local ec=0
    # run interactive job
    srun --cpus-per-task=$SJ_NCPUPROC -J "$SJ_JOBLABEL" --output="$stdout" --error="$stderr" $SJ_SCHEDOPTS < "$SJ_BATCH_SCRIPT"
    ec=$?
    rm -f "$SJ_BATCH_SCRIPT"
    if [[ $ec -ne 0 ]]; then
        error_msg "$SJ_JOBLABEL : slurm interactive job failed with error code $ec"
        return $ec
    fi
}

sj_accounting() {
    [[ $# -eq 0 ]] && return 0
    local jid=$1
    local cmd="sacct --format JobID,NodeList,JobName,MaxRSS,MinCPU -P -j $jid"
    local data; data=( `$cmd` )
    [[ $? -eq 0 ]] || { error_msg "# No accounting data - failed sacct -j $jid"; return 255; }

    local nlines=${#data[@]}
    [[ $nlines -ge 2 ]] || { error_msg "# No accounting data - produced no output sacct -j $jid"; return 255; }

    local numdata=
    local strdata=

    
    strdata=${data[1]}
    [[ $nlines -gt 2 ]] && numdata=${data[2]}

    echo "# Accounting obtained using: $cmd"
    echo "#JOBNAME $strdata"
    echo "#MEMORY  $numdata"
    echo "#CPUTIME $numdata"
}

sj_list_jobs() {
    [[ $# -gt 0 ]] || return 0
    local uid=$1
    local jobids=
    [[ $# -gt 1 ]] && jobids="--jobs=$2"
    #$CMD_SQUEUE -u $uid $jobids --noheader -o "%.7i" 2>/dev/null
    local cmd="$CMD_SQUEUE -u $uid $jobids --noheader -o %16i"
    debug_msg "list_jobs using : $cmd"
    #echoerr "DBG: command"
    #$cmd 1>&2
    #echoerr "DBG: done command"
    local res=0
    local ec=0
    res=$(retry_cmd_no_stderr 4 1 $cmd)
    ec=$?
    [[ $ec -eq 0 ]] || return $ec
    echo $res
    return 0
}

sj_cancel_jobid() {
    [[ $# -gt 0 ]] || return 0
    $CMD_SCANCEL $*
}
export -f sj_scheduler_init
export -f sj_submit_job
export -f sj_interactive_job
export -f sj_cancel_jobid
export -f sj_list_jobs
export -f sj_accounting

