#!/usr/bin/env bash
# alpha-rl-run.sh — One-command production training run.
#
# Convenience wrapper around argo-alpha-rl.sh that handles the full
# production flow: push, apply template, submit, monitor, download
# artifacts, and print a throughput summary.
#
# Usage:
#   ./scripts/alpha-rl-run.sh                    # defaults: 500k steps, b=256, L40S
#   ./scripts/alpha-rl-run.sh --steps 1000000    # 1M steps
#   ./scripts/alpha-rl-run.sh --batch 256        # b=256
#   ./scripts/alpha-rl-run.sh --gpu h100         # H100 pool
#   ./scripts/alpha-rl-run.sh --nsys             # with nsys profiling
#   ./scripts/alpha-rl-run.sh --dry-run          # print command, don't submit
#
# Flow:
#   1. Push current branch
#   2. Apply Argo template
#   3. Submit workflow
#   4. Monitor until completion
#   5. Download JSONL + summary from PVC
#   6. Print throughput + key metrics

set -euo pipefail

# ── Defaults ────────────────────────────────────────────────────────────
# Per feedback_default_to_l40s_pool: L40S is the default.
STEPS=500000
BATCH=256
GPU_POOL=ci-training-l40s
NSYS=false
BRANCH=$(git branch --show-current 2>/dev/null || echo "ml-alpha-phase-a")
DRY_RUN=false
SEED=16962
SEQ_LEN=32
PER_CAPACITY=32768
INSTRUMENT_MODE=all
FOLD_IDX=0
N_FOLDS=1
N_EVAL_STEPS=0

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

Production training run — push, submit, monitor, download, summarize.

Options:
  --steps <n>       Training steps (default: $STEPS)
  --batch <n>       Parallel LobSim instances per step (default: $BATCH)
  --gpu <pool>      l40s | h100 (default: l40s)
  --nsys            Enable nsys profiling
  --seed <n>        Random seed (default: $SEED)
  --seq-len <n>     Encoder sequence length (default: $SEQ_LEN)
  --per-capacity <n> PER buffer capacity (default: $PER_CAPACITY)
  --instrument <m>  all | front-month | id=<N> (default: $INSTRUMENT_MODE)
  --fold-idx <k>    Walk-forward fold index (default: $FOLD_IDX)
  --n-folds <K>     Total fold count (default: $N_FOLDS)
  --n-eval-steps <n> Eval-phase step count (default: $N_EVAL_STEPS)
  --branch <b>      Git branch (default: current branch)
  --dry-run         Print the submit command without executing
  -h, --help        Show this help
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --steps)         STEPS="$2"; shift 2 ;;
    --batch)         BATCH="$2"; shift 2 ;;
    --gpu)
      case "$2" in
        l40s)  GPU_POOL=ci-training-l40s ;;
        h100)  GPU_POOL=ci-training-h100 ;;
        *)     echo "ERROR: unknown gpu: $2 (use l40s or h100)"; exit 1 ;;
      esac
      shift 2 ;;
    --nsys)          NSYS=true; shift ;;
    --seed)          SEED="$2"; shift 2 ;;
    --seq-len)       SEQ_LEN="$2"; shift 2 ;;
    --per-capacity)  PER_CAPACITY="$2"; shift 2 ;;
    --instrument)    INSTRUMENT_MODE="$2"; shift 2 ;;
    --fold-idx)      FOLD_IDX="$2"; shift 2 ;;
    --n-folds)       N_FOLDS="$2"; shift 2 ;;
    --n-eval-steps)  N_EVAL_STEPS="$2"; shift 2 ;;
    --branch)        BRANCH="$2"; shift 2 ;;
    --dry-run)       DRY_RUN=true; shift ;;
    -h|--help)       usage; exit 0 ;;
    *) echo "Unknown option: $1"; usage; exit 1 ;;
  esac
done

# ── Step 1: Push current branch ────────────────────────────────────────
# Per feedback_push_before_deploy: always push before submitting.
echo "=== Step 1: Push $BRANCH ==="
if [[ "$DRY_RUN" == "false" ]]; then
  git push origin "$BRANCH"
  echo "  pushed: $BRANCH"
else
  echo "  [dry-run] would push: git push origin $BRANCH"
fi

SHA=$(git rev-parse --short=9 HEAD)
echo "  SHA: $SHA"

# ── Step 2: Apply Argo template ────────────────────────────────────────
# Per feedback_argo_template_must_apply: always apply before submit.
TEMPLATE_FILE="infra/k8s/argo/alpha-rl-template.yaml"
if [[ ! -f "$TEMPLATE_FILE" ]]; then
  echo "ERROR: template not found at $TEMPLATE_FILE (run from repo root)"
  exit 1
fi

echo ""
echo "=== Step 2: Apply Argo template ==="
if [[ "$DRY_RUN" == "false" ]]; then
  kubectl apply -n foxhunt -f "$TEMPLATE_FILE"
  echo "  applied: $TEMPLATE_FILE"
else
  echo "  [dry-run] would apply: kubectl apply -n foxhunt -f $TEMPLATE_FILE"
fi

# ── GPU pool → SM arch ──────────────────────────────────────────────────
case "$GPU_POOL" in
  ci-training-l40s)  SM=89 ;;
  ci-training-h100)  SM=90 ;;
esac

# ── Step 3: Submit workflow ─────────────────────────────────────────────
echo ""
echo "=== Step 3: Submit workflow ==="
echo "  branch:       $BRANCH"
echo "  sha:          $SHA"
echo "  gpu-pool:     $GPU_POOL (sm_$SM)"
echo "  steps:        $STEPS"
echo "  batch:        $BATCH"
echo "  seq-len:      $SEQ_LEN"
echo "  per-capacity: $PER_CAPACITY"
echo "  seed:         $SEED"
echo "  instrument:   $INSTRUMENT_MODE"
echo "  nsys:         $NSYS"

SUBMIT_CMD="argo submit -n foxhunt --from=wftmpl/alpha-rl \
  -p commit-sha=$SHA \
  -p git-branch=$BRANCH \
  -p cuda-compute-cap=$SM \
  -p gpu-pool=$GPU_POOL \
  -p n-steps=$STEPS \
  -p seq-len=$SEQ_LEN \
  -p n-backtests=$BATCH \
  -p per-capacity=$PER_CAPACITY \
  -p seed=$SEED \
  -p instrument-mode=$INSTRUMENT_MODE \
  -p fold-idx=$FOLD_IDX \
  -p n-folds=$N_FOLDS \
  -p n-eval-steps=$N_EVAL_STEPS \
  -p nsys-profile=$NSYS"

if [[ "$DRY_RUN" == "true" ]]; then
  echo ""
  echo "[dry-run] Would submit:"
  echo "  $SUBMIT_CMD"
  exit 0
fi

WF_OUTPUT=$(eval "$SUBMIT_CMD" 2>&1)
echo "$WF_OUTPUT"

# Extract workflow name from argo submit output
WF_NAME=$(echo "$WF_OUTPUT" | grep -oP 'Name:\s+\K\S+' || echo "")
if [[ -z "$WF_NAME" ]]; then
  echo "ERROR: could not extract workflow name from argo submit output"
  echo "       Check the output above and monitor manually."
  exit 1
fi
echo ""
echo "  workflow: $WF_NAME"

# ── Step 4: Monitor until completion ────────────────────────────────────
echo ""
echo "=== Step 4: Monitoring $WF_NAME ==="
echo "  (polling every 30s — Ctrl-C to detach, workflow continues)"
echo ""

START_TS=$(date +%s)

while true; do
  WF_PHASE=$(argo get -n foxhunt "$WF_NAME" -o json 2>/dev/null \
    | grep -oP '"phase"\s*:\s*"\K[^"]+' | head -1 || echo "Pending")

  case "$WF_PHASE" in
    Succeeded)
      echo ""
      echo "  SUCCEEDED after $(( ($(date +%s) - START_TS) / 60 ))m"
      break
      ;;
    Failed|Error)
      echo ""
      echo "  FAILED ($WF_PHASE) after $(( ($(date +%s) - START_TS) / 60 ))m"
      echo ""
      echo "=== Failure logs ==="
      argo logs -n foxhunt "$WF_NAME" --tail 50 2>/dev/null || true
      exit 1
      ;;
    *)
      printf "\r  status: %-12s elapsed: %dm" "$WF_PHASE" "$(( ($(date +%s) - START_TS) / 60 ))"
      sleep 30
      ;;
  esac
done

# ── Step 5: Download JSONL + summary from PVC ───────────────────────────
echo ""
echo "=== Step 5: Download artifacts ==="

# Find the pod that ran the workflow
POD=$(kubectl get pods -n foxhunt -l "workflows.argoproj.io/workflow=$WF_NAME" \
  --sort-by=.metadata.creationTimestamp -o jsonpath='{.items[-1].metadata.name}' 2>/dev/null || echo "")

OUT_DIR="/feature-cache/alpha-rl-runs/$SHA/fold$FOLD_IDX"
LOCAL_DIR="runs/$SHA"
mkdir -p "$LOCAL_DIR"

# Print final logs
echo ""
echo "=== Step 5a: Final training logs ==="
argo logs -n foxhunt "$WF_NAME" --tail 30 2>/dev/null \
  | grep -E "^step|complete|SUMMARY|throughput|sharpe|pnl" || true

# Copy summary JSON if the pod is still around
if [[ -n "$POD" ]]; then
  kubectl cp "foxhunt/$POD:$OUT_DIR/alpha_rl_train_summary.json" \
    "$LOCAL_DIR/alpha_rl_train_summary.json" 2>/dev/null || true
  kubectl cp "foxhunt/$POD:$OUT_DIR/diagnostics.jsonl" \
    "$LOCAL_DIR/diagnostics.jsonl" 2>/dev/null || true
  if [[ "$NSYS" == "true" ]]; then
    kubectl cp "foxhunt/$POD:$OUT_DIR/nsys_trace.nsys-rep" \
      "$LOCAL_DIR/nsys_trace.nsys-rep" 2>/dev/null || true
    kubectl cp "foxhunt/$POD:$OUT_DIR/nsys_trace.sqlite" \
      "$LOCAL_DIR/nsys_trace.sqlite" 2>/dev/null || true
  fi
  echo "  downloaded to: $LOCAL_DIR/"
  ls -lh "$LOCAL_DIR/" 2>/dev/null || true
else
  echo "  WARNING: pod not found — artifacts remain on PVC at $OUT_DIR"
  echo "           Use: kubectl cp foxhunt/<pod>:$OUT_DIR/ $LOCAL_DIR/"
fi

# ── Step 6: Print throughput + key metrics ──────────────────────────────
echo ""
echo "=== Step 6: Summary ==="
ELAPSED_MIN=$(( ($(date +%s) - START_TS) / 60 ))
echo "  wall time:  ${ELAPSED_MIN}m"
echo "  steps:      $STEPS"
echo "  batch:      $BATCH"
echo "  gpu:        $GPU_POOL"
echo "  sha:        $SHA"

if [[ -f "$LOCAL_DIR/alpha_rl_train_summary.json" ]]; then
  echo ""
  echo "  --- Training Summary ---"
  cat "$LOCAL_DIR/alpha_rl_train_summary.json"
  echo ""
fi

echo ""
echo "Done. Artifacts in: $LOCAL_DIR/"