#!/usr/bin/env bash
# Train a model via Argo Workflows (DQN, PPO, or any supervised model).
#
# Usage:
#   ./scripts/argo-train.sh dqn                           # DQN with defaults
#   ./scripts/argo-train.sh ppo --tag dev-abc1234          # PPO with specific binary
#   ./scripts/argo-train.sh tft --trials 40 --epochs 100   # supervised TFT
#   ./scripts/argo-train.sh mamba2 --webhook               # trigger via webhook
#
# Supported models:
#   RL:          dqn, ppo
#   Supervised:  tft, mamba2, tggn, tlob, liquid, kan, xlstm, diffusion
#
# Requires: argo CLI (default) or curl (webhook mode)
set -euo pipefail

TAG="latest"
TRIALS=""
EPOCHS=""
GPU_POOL=""
SYMBOL=""
WEBHOOK=false
WEBHOOK_BASE="http://workflow-trigger-eventsource-svc.foxhunt:12001"
WATCH=false

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

Models:
  dqn, ppo                                          (RL)
  tft, mamba2, tggn, tlob, liquid, kan, xlstm, diffusion  (supervised)

Options:
  --tag <binary-tag>     Binary tag to use (default: latest)
  --trials <n>           Hyperopt trials (default: model-specific)
  --epochs <n>           Training epochs (default: model-specific)
  --gpu-pool <pool>      GPU node pool (default: ci-training-h100)
  --symbol <sym>         Trading symbol (default: ES.FUT)
  --webhook              Trigger via webhook instead of argo CLI
  --watch                Follow workflow logs after submission
  -h, --help             Show this help
EOF
  exit 0
}

[[ $# -eq 0 ]] && { echo "Error: model argument required"; usage; }

MODEL="$1"; shift

# Validate model name
case "$MODEL" in
  dqn|ppo|tft|mamba2|tggn|tlob|liquid|kan|xlstm|diffusion) ;;
  *) echo "Error: unknown model '$MODEL'"; usage ;;
esac

while [[ $# -gt 0 ]]; do
  case $1 in
    --tag)       TAG="$2"; shift 2 ;;
    --trials)    TRIALS="$2"; shift 2 ;;
    --epochs)    EPOCHS="$2"; shift 2 ;;
    --gpu-pool)  GPU_POOL="$2"; shift 2 ;;
    --symbol)    SYMBOL="$2"; shift 2 ;;
    --webhook)   WEBHOOK=true; shift ;;
    --watch)     WATCH=true; shift ;;
    -h|--help)   usage ;;
    *)           echo "Unknown option: $1"; usage ;;
  esac
done

# Determine workflow template name
case "$MODEL" in
  dqn)  TEMPLATE="train-dqn" ;;
  ppo)  TEMPLATE="train-ppo" ;;
  *)    TEMPLATE="train-supervised" ;;
esac

if $WEBHOOK; then
  case "$MODEL" in
    dqn) ENDPOINT="/train/dqn";        PAYLOAD="{\"binary_tag\": \"${TAG}\"}" ;;
    ppo) ENDPOINT="/train/ppo";        PAYLOAD="{\"binary_tag\": \"${TAG}\"}" ;;
    *)   ENDPOINT="/train/supervised"; PAYLOAD="{\"model\": \"${MODEL}\", \"binary_tag\": \"${TAG}\"}" ;;
  esac

  echo "Triggering $MODEL training via webhook..."
  curl -sf -X POST "${WEBHOOK_BASE}${ENDPOINT}" \
    -H 'Content-Type: application/json' \
    -d "$PAYLOAD"
  echo "Webhook sent: $PAYLOAD"
  exit 0
fi

# Build argo submit command
CMD="argo submit -n foxhunt --from=wftmpl/$TEMPLATE"
CMD="$CMD -p binary-tag=$TAG"

# Supervised models need --model param
if [[ "$MODEL" != "dqn" && "$MODEL" != "ppo" ]]; then
  CMD="$CMD -p model=$MODEL"
fi

[[ -n "$TRIALS" ]]   && CMD="$CMD -p hyperopt-trials=$TRIALS"
[[ -n "$EPOCHS" ]]   && CMD="$CMD -p train-epochs=$EPOCHS"
[[ -n "$GPU_POOL" ]] && CMD="$CMD -p gpu-pool=$GPU_POOL"
[[ -n "$SYMBOL" ]]   && CMD="$CMD -p symbol=$SYMBOL"

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

echo "Submitting $MODEL training workflow ($TEMPLATE)..."
echo "  binary-tag: $TAG"
[[ -n "$TRIALS" ]]   && echo "  trials:     $TRIALS"
[[ -n "$EPOCHS" ]]   && echo "  epochs:     $EPOCHS"
[[ -n "$GPU_POOL" ]] && echo "  gpu-pool:   $GPU_POOL"
[[ -n "$SYMBOL" ]]   && echo "  symbol:     $SYMBOL"
echo ""
eval "$CMD"