##
##  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
    ## handles GNU coreutils (9-digit %N), uutils coreutils (zero-stripped %N),
    ## and BSD/macOS date (literal "N" from %N, falling back to second precision)
    ## note: seconds and nanoseconds are read with a single `date` invocation so
    ## they are sampled atomically and cannot straddle a second boundary
    local s n
    read s n < <(date '+%s %N')
    if [[ $n =~ ^[0-9]+$ ]]; then
        n=$(printf '%09d' "$((10#$n))")
        echo "$((s * 1000 + 10#${n:0:3}))"
    else
        echo "$((s * 1000))"
    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() {
    ## This is invoked via `trap on_exit EXIT` at the end of nxf_main, therefore
    ## $? here is the exit status of the very last command run by nxf_main, i.e.
    ## nxf_unstage_controls (or {{after_script}}, if the task defines one).
    ## It's kept as a last-resort fallback ONLY: nxf_main_ret and nxf_unstage_ret
    ## below already capture the two "expected" failure points (the task command
    ## and the output unstaging step), so last_err only wins when the failure
    ## happened somewhere else, e.g. in nxf_unstage_controls, in the after script,
    ## or the trap fired directly because of `set -e`/`set -u` without either
    ## variable being set.
    local last_err=$?
    ## Exit status precedence (first non-zero wins):
    ##   1) nxf_main_ret    - exit status of the user task script, set by
    ##                        `wait $pid || nxf_main_ret=$?` once nxf_launch completes.
    ##   2) nxf_unstage_ret - exit status of the nxf_unstage_outputs pipeline, set in
    ##                        nxf_unstage() when copying declared outputs back fails.
    ##                        Only meaningful when nxf_main_ret==0, since outputs are
    ##                        not unstaged if the task itself already failed.
    ##   3) last_err        - fallback for anything else, see above.
    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() {
    ## Only unstage the task outputs if the task command itself succeeded;
    ## if nxf_main_ret is non-zero, there is nothing new to copy back and
    ## nxf_unstage_ret is left unset (defaults to 0) so it can't mask the
    ## task's own error in on_exit.
    if [[ ${nxf_main_ret:=0} == 0 ]]; then
        ## Run the outputs-unstaging script in its own subshell with
        ## `set -e -o pipefail` (fast failure is off in the outer script at
        ## this point, see nxf_main) so that a failed copy/upload is caught
        ## via `nxf_unstage_ret=$?` below instead of silently continuing.
        ## stdout/stderr of the subshell are appended to .command.out/.err
        ## (rather than replacing them) so any output produced while
        ## unstaging is visible next to the task's own output.
        (set -e -o pipefail; (nxf_unstage_outputs | tee -a {{stdout_file}}) 3>&1 1>&2 2>&3 | tee -a {{stderr_file}}) 3>&1 1>&2 2>&3
        nxf_unstage_ret=$?
    fi
    ## Control files (.command.out/.err/.log/...) are unstaged unconditionally,
    ## even if the task or the outputs unstaging above failed, so the logs
    ## needed to diagnose the failure are always copied back. Its own exit
    ## status is intentionally NOT captured into nxf_unstage_ret: it becomes
    ## the generic `last_err` fallback in on_exit instead (see there), unless
    ## an after script runs later, whose exit status then takes its place.
    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}}) 3>&1 1>&2 2>&3 &
    pid=$!
    wait $pid || nxf_main_ret=$?
    {{unstage_cmd}}
    {{after_script}}
}

$NXF_ENTRY
