#!/bin/bash
# Usage: ./run_and_wait.sh <yaml-file> [poll-interval-seconds]
# Deletes existing job, applies yaml, waits for build+run, prints results.

set -euo pipefail

YAML="${1:?Usage: $0 <yaml-file> [poll-interval]}"
POLL="${2:-10}"

JOB_NAME=$(grep -m1 '^\s*name:' "$YAML" | awk '{print $2}')
NAMESPACE=$(grep -m1 '^\s*namespace:' "$YAML" | awk '{print $2}')
LABEL="batch.kubernetes.io/job-name=${JOB_NAME}-node-0"

echo "=== $JOB_NAME | $NAMESPACE ==="

# Remember old pods to ignore them
OLD_PODS=$(kubectl get pods -n "$NAMESPACE" -l "$LABEL" -o jsonpath='{.items[*].metadata.name}' 2>/dev/null || true)

# Delete if exists
kubectl delete -f "$YAML" 2>/dev/null && {
  echo "Deleted old job. Waiting for cleanup..."
  while kubectl get trainjob -n "$NAMESPACE" "$JOB_NAME" &>/dev/null; do sleep 2; done
  sleep 3
}

# Apply
kubectl apply -f "$YAML"

# Wait for NEW pod (not in OLD_PODS)
echo -n "Waiting for new pod"
POD=""
while true; do
  ALL_PODS=$(kubectl get pods -n "$NAMESPACE" -l "$LABEL" -o jsonpath='{.items[*].metadata.name}' 2>/dev/null || true)
  for p in $ALL_PODS; do
    if [[ ! " $OLD_PODS " =~ " $p " ]]; then
      PHASE=$(kubectl get pod -n "$NAMESPACE" "$p" -o jsonpath='{.status.phase}' 2>/dev/null || true)
      if [[ -n "$PHASE" ]]; then
        POD="$p"
        echo " $POD ($PHASE)"
        break 2
      fi
    fi
  done
  echo -n "."
  sleep "$POLL"
done

# Wait for pod to complete
echo -n "Running"
while true; do
  PHASE=$(kubectl get pod -n "$NAMESPACE" "$POD" -o jsonpath='{.status.phase}' 2>/dev/null || echo "Gone")
  if [[ "$PHASE" != "Running" && "$PHASE" != "Pending" ]]; then
    echo " ($PHASE)"
    break
  fi
  echo -n "."
  sleep "$POLL"
done

# Print logs
echo ""
echo "=== LOGS ==="
kubectl logs -n "$NAMESPACE" "$POD" 2>/dev/null || echo "(no logs available)"

echo ""
echo "=== STATUS: $PHASE ==="