#!/usr/bin/env bash
# Run GPU/CUDA tests via Argo Workflows on H100.
#
# Usage:
#   ./scripts/argo-test.sh                              # defaults: dqn,ppo,tft
#   ./scripts/argo-test.sh --models dqn --scope lib     # DQN lib tests only
#   ./scripts/argo-test.sh --models all                  # all 10 models
#   ./scripts/argo-test.sh --watch                      # follow logs
#
# Requires: argo CLI
set -euo pipefail

MODELS="dqn,ppo,tft"
SCOPE="all"
COMMIT_REF="HEAD"
GPU_POOL=""
WATCH=false

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

Options:
  --models <list>      Comma-separated model list (default: dqn,ppo,tft)
                       Use "all" for all 10 models
  --scope <scope>      lib, integration, or all (default: all)
  --ref <ref>          Git ref to test (default: HEAD)
  --gpu-pool <pool>    GPU node pool (default: ci-training-h100)
  --watch              Follow workflow logs
  -h, --help           Show this help
EOF
  exit 0
}

while [[ $# -gt 0 ]]; do
  case $1 in
    --models)    MODELS="$2"; shift 2 ;;
    --scope)     SCOPE="$2"; shift 2 ;;
    --ref)       COMMIT_REF="$2"; shift 2 ;;
    --gpu-pool)  GPU_POOL="$2"; shift 2 ;;
    --watch)     WATCH=true; shift ;;
    -h|--help)   usage ;;
    *)           echo "Unknown option: $1"; usage ;;
  esac
done

# Resolve HEAD to actual SHA if needed
if [ "$COMMIT_REF" = "HEAD" ]; then
  COMMIT_REF=$(git rev-parse HEAD)
fi

CMD="argo submit -n foxhunt --from=wftmpl/gpu-test-pipeline"
CMD="$CMD -p commit-ref=$COMMIT_REF"
CMD="$CMD -p models=$MODELS"
CMD="$CMD -p test-scope=$SCOPE"

[[ -n "$GPU_POOL" ]] && CMD="$CMD -p gpu-pool=$GPU_POOL"

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

echo "Submitting GPU test workflow..."
echo "  commit:  $(echo $COMMIT_REF | cut -c1-8)"
echo "  models:  $MODELS"
echo "  scope:   $SCOPE"
echo ""
eval "$CMD"