#!/bin/bash
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# Launcher script to avoid early-exit preventing output from being printed
PROJECT_DIR="$1"
TEST_TYPE="$2"

SBATCH_LOG="$PROJECT_DIR/nvshmem4py_sbatch.out"
TEST_LOG="$PROJECT_DIR/nvshmem4py_test.out"

source $PROJECT_DIR/docker/clusters/eos-config.sh
configure_test_env
CONTAINER_IMAGE=$(get_build_tools_image $(get_build_image_version))

date

# launch sbatch
# SBATCH_EXTRA_ARGS can be set by CI jobs to override sbatch defaults (e.g. node count)
# Pass the cluster-selected runtime environment into the container.  The EOS
# config supplies CUDA/OpenMPI/NCCL paths here; dropping them makes NVLS tests
# pick up incompatible library fallbacks from the Python environment.
CONTAINER_ENV_NAMES=()
for var in \
    CUDA_HOME MPI_HOME LD_LIBRARY_PATH PATH \
    NVSHMEM_HOME NVSHMEM_PREFIX RDMA_CORE_HOME \
    NVSHMEM_ENABLE_HCA_PE_MAPPING NVSHMEM_HCA_LIST NVSHMEM_MAX_TEAMS \
    NVSHMEM_DEBUG NVSHMEM_BOOTSTRAP NP NP_TEAM; do
    if [ -n "${!var+x}" ]; then
        CONTAINER_ENV_NAMES+=("$var")
    fi
done
CONTAINER_ENV_ARGS=()
if [ ${#CONTAINER_ENV_NAMES[@]} -ne 0 ]; then
    IFS=,
    CONTAINER_ENV_ARGS=(--container-env="${CONTAINER_ENV_NAMES[*]}")
    unset IFS
fi
echo "Forwarding container env: ${CONTAINER_ENV_NAMES[*]:-<none>}"
sbatch -W ${SBATCH_EXTRA_ARGS:-} --container-image=$CONTAINER_IMAGE --container-mounts=$PROJECT_DIR/ \
    "${CONTAINER_ENV_ARGS[@]}" \
    $PROJECT_DIR/nvshmem4py/test/python_test_wrapper.sh $PROJECT_DIR $TEST_TYPE > $SBATCH_LOG 2>&1 &
SBATCH_PID=$!

# TODO(vkhodel): Matrix uses a loop to wait for job id info
sleep 5

JOB_ID=$(cat $SBATCH_LOG | sed -e 's/^.* \([0-9]\+\)$/\1/')
echo JOB_ID is $JOB_ID

# TODO(vkhodel): CodeRabbit suggests to use a loop here to wait for Pending state
squeue -j $JOB_ID --start

# show output, wait for the file to appear, exit on sbatch exit
tail --retry --sleep-interval 5 --pid $SBATCH_PID -f $TEST_LOG &

wait $SBATCH_PID
SBATCH_EXIT_CODE=$?

date

# this might fail, let's skip for now
#EXIT_CODE=$(sacct -j "$JOB_ID" --format=ExitCode --noheader | awk -F: '{print $1}' | head -1)
#
#if [ -z "$EXIT_CODE" ]; then
#    echo "ERROR: No exit code for job $JOB_ID found"
#    exit 1
#fi
#
#echo "Slurm job exited with $EXIT_CODE"

echo "Sbatch exited with $SBATCH_EXIT_CODE"
exit $SBATCH_EXIT_CODE
