#!/bin/bash
set -euo pipefail

shopt -s nullglob

# -------------------------
# parse command line arguments
# -------------------------
SUFFIX=""
DO_TEST=0
DO_NSYS=0
DO_NCU=0

while [[ $# -gt 0 ]]; do
    case $1 in
        --test)
            DO_TEST=1
            shift
            ;;
        --nsys)
            DO_NSYS=1
            shift
            ;;
        --ncu)
            DO_NCU=1
            shift
            ;;
        *)
            if [[ -z "$SUFFIX" ]]; then
                SUFFIX="$1"
            else
                echo "Error: multiple non-option arguments provided ('$SUFFIX' and '$1')"
                exit 1
            fi
            shift
            ;;
    esac
done

if [[ -z "$SUFFIX" ]]; then
    echo "Usage: $0 <suffix> [--test] [--nsys] [--ncu]"
    exit 1
fi

echo "Suffix: $SUFFIX"

# -------------------------
# kernel swap (temporary replacement)
# -------------------------
SRC_DIR="kernels_${SUFFIX}"
KERNEL_DIR="kernels"

if [[ ! -d "$SRC_DIR" ]]; then
    echo "Error: source directory '$SRC_DIR' does not exist."
    exit 1
fi

src_files=("$SRC_DIR"/*.cuh)
if [[ ${#src_files[@]} -eq 0 ]]; then
    echo "Error: no .cuh files found in '$SRC_DIR'"
    exit 1
fi

BACKUP_DIR=".kernel_backup_${SUFFIX}_$(date +%s)"
mkdir -p "$BACKUP_DIR"

backup_files=("$KERNEL_DIR"/*.cuh)
if [[ ${#backup_files[@]} -gt 0 ]]; then
    echo "Backing up current kernel files to $BACKUP_DIR"
    cp "${backup_files[@]}" "$BACKUP_DIR/"
else
    echo "No existing .cuh files in $KERNEL_DIR to back up"
fi

# -------------------------
# trap: always restore on exit (success, error, or signal)
# -------------------------
trap '
    echo "Restoring original kernel files from backup..."
    restore_files=("$BACKUP_DIR"/*.cuh)
    if [[ ${#restore_files[@]} -gt 0 ]]; then
        for f in "${restore_files[@]}"; do
            cp "$f" "$KERNEL_DIR/"
        done
        echo "Restoration complete."
    else
        echo "Warning: no backup files found in $BACKUP_DIR — kernel directory may contain swapped files."
    fi
    rm -rf "$BACKUP_DIR"
' EXIT

echo "Swapping kernel files from $SRC_DIR to $KERNEL_DIR"
for f in "${src_files[@]}"; do
    base=$(basename "$f")
    cp "$f" "$KERNEL_DIR/$base"
done

# -------------------------
# build the executable
# -------------------------
echo "Building test_gpt2_${SUFFIX}..."
make test_gpt2 "$SUFFIX"  # if this fails, trap fires and restores

BIN="./build/test_gpt2_${SUFFIX}"

# =========================================================
# TEST RUN (correctness)
# =========================================================
if [[ $DO_TEST -eq 1 ]]; then
    echo "Running TEST (sbatch verify.slurm $SUFFIX)..."
    sbatch verify.slurm "$SUFFIX"
fi

# =========================================================
# NSYS RUN (profiling with Nsight Systems)
# =========================================================
if [[ $DO_NSYS -eq 1 ]]; then
    echo "Running NSYS (sbatch verify_nsys.slurm $SUFFIX)..."
    sbatch verify_nsys.slurm "$SUFFIX"
fi

# =========================================================
# NCU RUN (profiling with Nsight Compute)
# =========================================================
if [[ $DO_NCU -eq 1 ]]; then
    echo "Running NCU (sbatch verify_ncu.slurm $SUFFIX)..."
    sbatch verify_ncu.slurm "$SUFFIX"
fi

# -------------------------
# final status
# -------------------------
echo "Checking Slurm queue for user $USER..."
squeue -u "$USER"

echo "Done."
# trap fires here on clean exit too — restore always runs