#!/usr/bin/env bash
# Run a smoke test under compute-sanitizer on L40S via Argo Workflows.
#
# Why L40S not local: 4 GB laptop GPU can't fit memcheck instrumentation
# overhead alongside production training kernels (sanitizer falls back to
# "didn't track the launch"). L40S 48 GB has ample headroom.
#
# Usage:
#   ./scripts/argo-sanitizer.sh                                          # iqn_quantile_monotonicity, memcheck
#   ./scripts/argo-sanitizer.sh --test magnitude_distribution            # different test
#   ./scripts/argo-sanitizer.sh --tool racecheck                         # different sanitizer tool
#   ./scripts/argo-sanitizer.sh --watch                                  # follow logs
#   ./scripts/argo-sanitizer.sh --ref main --tool memcheck --watch
#
# Tool options: memcheck (OOB/leaks), racecheck (shared-mem races),
#               synccheck (sync API misuse), initcheck (uninitialised reads)
#
# Requires: argo CLI

set -euo pipefail

TEST_NAME="iqn_quantile_monotonicity::iqn_multi_quantile_heads_produce_monotonic_estimates"
SANITIZER_TOOL="memcheck"
COMMIT_REF="HEAD"
GPU_POOL=""
EXTRA_ARGS=""
WATCH=false

usage() {
  cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Wraps a smoke test under compute-sanitizer on L40S.

Options:
  --test <name>        Smoke test name (default: $TEST_NAME)
                       Common values:
                         iqn_quantile_monotonicity::iqn_multi_quantile_heads_produce_monotonic_estimates
                         multi_fold_convergence::test_multi_fold_convergence
                         magnitude_distribution::test_magnitude_distribution
                         performance::test_real_data_single_epoch
  --multi-fold         Shortcut: --test multi_fold_convergence::test_multi_fold_convergence
                       Exercises fold-boundary code paths (IQN sync, aux Adam reset,
                       iqn_readiness reset, MSE clamp) that single-fold tests cannot.
                       Walk-forward needs ≥10 months of data — uses CI test-data PVC.
  --tool <tool>        Sanitizer tool (default: memcheck)
                       memcheck | racecheck | synccheck | initcheck
  --ref <ref>          Git ref to test (default: HEAD)
  --gpu-pool <pool>    GPU node pool (default: ci-training-l40s)
  --extra-args <args>  Pass-through args to compute-sanitizer
                       e.g. "--leak-check full" for memcheck
  --watch              Follow workflow logs
  -h, --help           Show this help
EOF
  exit 0
}

while [[ $# -gt 0 ]]; do
  case $1 in
    --test)        TEST_NAME="$2"; shift 2 ;;
    --multi-fold)  TEST_NAME="multi_fold_convergence::test_multi_fold_convergence"; shift ;;
    --tool)        SANITIZER_TOOL="$2"; shift 2 ;;
    --ref)         COMMIT_REF="$2"; shift 2 ;;
    --gpu-pool)    GPU_POOL="$2"; shift 2 ;;
    --extra-args)  EXTRA_ARGS="$2"; shift 2 ;;
    --watch)       WATCH=true; shift ;;
    -h|--help)     usage ;;
    *)             echo "Unknown option: $1"; usage ;;
  esac
done

# Resolve HEAD to actual SHA so the workflow targets a stable commit
if [ "$COMMIT_REF" = "HEAD" ]; then
  COMMIT_REF=$(git rev-parse HEAD)
fi

# Validate sanitizer tool
case "$SANITIZER_TOOL" in
  memcheck|racecheck|synccheck|initcheck) ;;
  *)
    echo "ERROR: invalid --tool '$SANITIZER_TOOL'"
    echo "Allowed: memcheck | racecheck | synccheck | initcheck"
    exit 1
    ;;
esac

CMD="argo submit -n foxhunt --from=wftmpl/sanitizer-test"
CMD="$CMD -p commit-ref=$COMMIT_REF"
CMD="$CMD -p test-name=$TEST_NAME"
CMD="$CMD -p sanitizer-tool=$SANITIZER_TOOL"
[[ -n "$GPU_POOL" ]]      && CMD="$CMD -p gpu-pool=$GPU_POOL"
[[ -n "$EXTRA_ARGS" ]]    && CMD="$CMD -p sanitizer-extra-args=\"$EXTRA_ARGS\""

if $WATCH; then
  CMD="$CMD --watch"
fi

echo "Submitting compute-sanitizer workflow..."
echo "  commit: $(echo "$COMMIT_REF" | cut -c1-8)"
echo "  test:   $TEST_NAME"
echo "  tool:   $SANITIZER_TOOL"
[[ -n "$EXTRA_ARGS" ]] && echo "  extra:  $EXTRA_ARGS"
echo ""
eval "$CMD"