#!/bin/bash --
# chaperon/stubs/sbatch — PATH-shadowing stub for sbatch
#
# Sends sbatch requests to the chaperon process via named pipes (_CHAPERON_FIFO_DIR).
# Runs INSIDE the sandbox; the chaperon validates, wraps, and submits
# the job on the outside.
#
# NOTE: Unlike all other stubs (which are trivial `chaperon_call <cmd> "$@"`
# passthroughs), this stub does client-side argument parsing.  This is
# intentional: sbatch takes a script *file* as a positional argument, and that
# file may live on a sandbox tmpfs that the handler (running outside) cannot read.
# So the stub must extract the script content here and send it via the SCRIPT
# protocol field.  The handler then only sees validated flags + inline script.
# srun has no script-file argument, so its stub needs no parsing.

set -euo pipefail

_STUB_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
source "$_STUB_DIR/_stub_lib.sh"

# ── Parse arguments to find --wrap and script file ───────────────
# We need to intercept --wrap (send as SCRIPT) and the script file
# (read and send as SCRIPT). Everything else is forwarded as ARGs.

FORWARD_ARGS=()
WRAP_CMD=""
SCRIPT_FILE=""
SCRIPT_ARGS=()

i=0
args=("$@")
while (( i < ${#args[@]} )); do
    arg="${args[$i]}"
    case "$arg" in
        --wrap=*)
            WRAP_CMD="${arg#--wrap=}"
            ;;
        --wrap)
            i=$(( i + 1 ))
            WRAP_CMD="${args[$i]:-}"
            ;;
        *)
            FORWARD_ARGS+=("$arg")
            ;;
    esac
    (( i++ )) || true
done

# If --wrap was given, convert to a script
if [[ -n "$WRAP_CMD" ]]; then
    _CHAPERON_SCRIPT="#!/bin/bash --
$WRAP_CMD"
    export _CHAPERON_SCRIPT
    chaperon_call sbatch "${FORWARD_ARGS[@]}"
    exit $?
fi

# No --wrap: look for a script file among the forward args.
# Strategy: walk args with flag-value awareness to find the first
# bare positional that is a readable file.
#
# We use a known list of flags that consume a value argument so we
# can correctly distinguish "flag value" pairs from positional args
# (the script file).  Without this, boolean flags like --parsable
# cause the next positional to be mistaken for a flag value.

# Flags that consume a subsequent value token (space-separated form).
# Keep in sync with _SBATCH_VALUE_FLAGS in handlers/_handler_lib.sh.
_STUB_VALUE_FLAGS=" \
  -A --account \
  -c --cpus-per-task \
  -d --dependency \
  -e --error \
  -J --job-name \
  -n --ntasks \
  -N --nodes \
  -o --output \
  -p --partition \
  -q --qos \
  -t --time \
  -G --gpus \
  -w --nodelist \
  -x --exclude \
  --begin \
  --comment \
  --constraint \
  --cpu-freq \
  --deadline \
  --gres \
  --gpus-per-node \
  --gpus-per-task \
  --mail-type \
  --mail-user \
  --mem \
  --mem-per-cpu \
  --mem-per-gpu \
  --nice \
  --ntasks-per-node \
  --cpus-per-gpu \
  --priority \
  --reservation \
  --signal \
  --switches \
  --threads-per-core \
  --tmp \
  --wait-all-nodes \
  --wckey \
  --array \
"

FINAL_ARGS=()
skip_next=false

for (( i=0; i < ${#FORWARD_ARGS[@]}; i++ )); do
    arg="${FORWARD_ARGS[$i]}"

    if "$skip_next"; then
        FINAL_ARGS+=("$arg")
        skip_next=false
        continue
    fi

    case "$arg" in
        --*=*)
            FINAL_ARGS+=("$arg")
            ;;
        -*)
            FINAL_ARGS+=("$arg")
            # Only skip the next token if this flag takes a value
            if [[ "$_STUB_VALUE_FLAGS" == *" $arg "* ]] \
               && (( i + 1 < ${#FORWARD_ARGS[@]} )); then
                skip_next=true
            fi
            ;;
        *)
            if [[ -z "$SCRIPT_FILE" && -f "$arg" ]]; then
                SCRIPT_FILE="$arg"
                # Everything after the script file is script args
                SCRIPT_ARGS=("${FORWARD_ARGS[@]:$((i+1))}")
                break
            else
                FINAL_ARGS+=("$arg")
            fi
            ;;
    esac
done

if [[ -n "$SCRIPT_FILE" ]]; then
    # Read the script file and send as SCRIPT.
    # Forward any positional args after the script file as SCRIPT_ARGs so
    # they reach $1/$@/$# inside the user's wrapped script on the compute node.
    _CHAPERON_SCRIPT="$(cat "$SCRIPT_FILE")"
    export _CHAPERON_SCRIPT
    _CHAPERON_SCRIPT_ARGS=("${SCRIPT_ARGS[@]+"${SCRIPT_ARGS[@]}"}")
    chaperon_call sbatch "${FINAL_ARGS[@]}"
    exit $?
fi

# No --wrap and no script file — pass through (e.g., --help, --version)
chaperon_call sbatch "${FORWARD_ARGS[@]}"
exit $?
