#!/bin/bash
# Usage: ./run_bench.sh <yaml-file> [poll-interval-seconds]
# Example: ./run_bench.sh benchmark_rms_optim.yaml 10

set -euo pipefail

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

if [[ ! -f "$YAML" ]]; then
  echo "Error: $YAML not found"
  exit 1
fi

# Extract job name and namespace from yaml
JOB_NAME=$(grep -m1 '^\s*name:' "$YAML" | awk '{print $2}')
NAMESPACE=$(grep -m1 '^\s*namespace:' "$YAML" | awk '{print $2}')

echo "=== Job: $JOB_NAME | Namespace: $NAMESPACE ==="

# Delete if exists
kubectl delete trainjob -n "$NAMESPACE" "$JOB_NAME" 2>/dev/null && {
  echo "Deleted existing job, waiting 5s..."
  sleep 5
}

# Apply
kubectl apply -f "$YAML"
echo "Applied. Polling every ${POLL}s..."

# Wait for pod
echo -n "Waiting for pod..."
while true; do
  POD=$(kubectl get pods -n "$NAMESPACE" -l "batch.kubernetes.io/job-name=${JOB_NAME}-node-0" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)
  if [[ -n "$POD" && "$POD" != "" ]]; then
    STATUS=$(kubectl get pod -n "$NAMESPACE" "$POD" -o jsonpath='{.status.phase}' 2>/dev/null || true)
    if [[ "$STATUS" == "Running" || "$STATUS" == "Succeeded" || "$STATUS" == "Failed" ]]; then
      echo " $POD ($STATUS)"
      break
    fi
  fi
  echo -n "."
  sleep "$POLL"
done

# Stream logs until completion
echo "=== Streaming logs ==="
kubectl logs -n "$NAMESPACE" "$POD" -f 2>/dev/null || true

# Final status
echo ""
echo "=== Final Status ==="
kubectl get trainjob -n "$NAMESPACE" "$JOB_NAME" -o jsonpath='{.status.conditions[0].reason}' 2>/dev/null
echo ""