##
##  Copyright 2013-2024, Seqera Labs
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##      http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.
#!/bin/bash
{{header_script}}
{{task_metadata}}
set -e
set -u
NXF_DEBUG=${NXF_DEBUG:=0}; [[ $NXF_DEBUG > 1 ]] && set -x
NXF_ENTRY=${1:-nxf_main}

{{trace_script}}
{{helpers_script}}
{{container_env}}

nxf_sleep() {
  sleep $1 2>/dev/null || sleep 1;
}

nxf_date() {
    ## should return the current timestamp in milliseconds
    ## note1: some linux silently ignores the `%3N` option and returns the ts in seconds (len==10)
    ## note2: old date tool ignores the `%3N` option and append that string to the timestamp in seconds
    ## note3: mac date tools ignores the `%3N` option and the string `3N` is appended to the timestamp in seconds
    local ts=$(date +%s%3N);
    if [[ ${#ts} == 10 ]]; then echo ${ts}000
    elif [[ $ts == *%3N ]]; then echo ${ts/\%3N/000}
    elif [[ $ts == *3N ]]; then echo ${ts/3N/000}
    elif [[ ${#ts} == 13 ]]; then echo $ts
    else echo "Unexpected timestamp value: $ts"; exit 1
    fi
}

nxf_env() {
    echo '============= task environment ============='
    env | sort | sed "s/\(.*\)AWS\(.*\)=\(.\{6\}\).*/\1AWS\2=\3xxxxxxxxxxxxx/"
    echo '============= task output =================='
}

nxf_kill() {
    declare -a children
    while read P PP;do
        children[$PP]+=" $P"
    done < <(ps -e -o pid= -o ppid=)

    kill_all() {
        [[ $1 != $$ ]] && kill $1 2>/dev/null || true
        for i in ${children[$1]:=}; do kill_all $i; done
    }

    kill_all $1
}

nxf_mktemp() {
    local base=${1:-/tmp}
    mkdir -p "$base"
    if [[ $(uname) = Darwin ]]; then mktemp -d $base/nxf.XXXXXXXXXX
    else TMPDIR="$base" mktemp -d -t nxf.XXXXXXXXXX
    fi
}

nxf_fs_copy() {
  local source=$1
  local target=$2
  local basedir=$(dirname $1)
  mkdir -p $target/$basedir
  cp -fRL $source $target/$basedir
}

nxf_fs_move() {
  local source=$1
  local target=$2
  local basedir=$(dirname $1)
  mkdir -p $target/$basedir
  mv -f $source $target/$basedir
}

nxf_fs_rsync() {
  local source=$1
  local target=$2
  mkdir -p $target
  rsync -rRl $source $target
}

nxf_fs_rclone() {
  rclone copyto $1 $2/$1
}

nxf_fs_fcp() {
  fcp $1 $2/$1
}

on_exit() {
    ## Capture possible errors.
    ## Can be caused either by the task script, unstage script or after script if defined
    local last_err=$?
    ## capture the task error first or fallback to unstage error
    local exit_status=${nxf_main_ret:=0}
    [[ ${exit_status} -eq 0 && ${nxf_unstage_ret:=0} -ne 0 ]] && exit_status=${nxf_unstage_ret:=0}
    [[ ${exit_status} -eq 0 && ${last_err} -ne 0 ]] && exit_status=${last_err}
    printf -- $exit_status {{exit_file}}
    set +u
    {{cleanup_cmd}}
    {{sync_cmd}}
    exit $exit_status
}

on_term() {
    set +e
    {{kill_cmd}}
}

nxf_launch() {
    {{launch_cmd}}
}

nxf_stage() {
    true
    {{stage_inputs}}
}

nxf_unstage_outputs() {
    true
    {{unstage_outputs}}
}

nxf_unstage_controls() {
    true
    {{unstage_controls}}
}

nxf_unstage() {
    ## Deactivate fast failure to allow uploading stdout and stderr files later
    if [[ ${nxf_main_ret:=0} == 0 ]]; then
        ## Data unstaging redirecting stdout and stderr with append mode
        (set -e -o pipefail; (nxf_unstage_outputs | tee -a {{stdout_file}}) 3>&1 1>&2 2>&3 | tee -a {{stderr_file}})
        nxf_unstage_ret=$?
    fi
    nxf_unstage_controls
}

nxf_main() {
    trap on_exit EXIT
    trap on_term TERM INT USR2
    trap '' USR1

    [[ "${NXF_CHDIR:-}" ]] && cd "$NXF_CHDIR"
    {{container_boxid}}
    {{scratch_cmd}}
    [[ $NXF_DEBUG > 0 ]] && nxf_env
    {{touch_file}}
    set +u
    {{before_script}}
    {{module_load}}
    {{conda_activate}}
    {{spack_activate}}
    set -u
    {{task_env}}
    {{secrets_env}}
    [[ $NXF_SCRATCH ]] && cd $NXF_SCRATCH
    export NXF_TASK_WORKDIR="$PWD"
    {{stage_cmd}}

    set +e
    (set -o pipefail; (nxf_launch | tee {{stdout_file}}) 3>&1 1>&2 2>&3 | tee {{stderr_file}}) &
    pid=$!
    wait $pid || nxf_main_ret=$?
    {{unstage_cmd}}
    {{after_script}}
}

$NXF_ENTRY
