#!/bin/bash

# Runpod Pod Monitoring Script
# Usage: ./scripts/monitor_runpod.sh [pod_id]
# Default pod: six9yaydb5f19b

set -e

# Configuration
POD_ID="${1:-six9yaydb5f19b}"
CHECK_INTERVAL=30  # seconds between checks
MAX_WAIT=900       # 15 minutes max wait time

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

echo "========================================="
echo "Runpod Pod Monitoring Script"
echo "========================================="
echo "Pod ID: $POD_ID"
echo "Check Interval: ${CHECK_INTERVAL}s"
echo "Max Wait Time: ${MAX_WAIT}s ($(($MAX_WAIT / 60)) minutes)"
echo "========================================="
echo ""

# Function to check pod status
check_pod_status() {
    runpodctl get pod "$POD_ID" 2>&1
}

# Function to check SSH readiness
check_ssh_ready() {
    OUTPUT=$(runpodctl ssh connect "$POD_ID" 2>&1)
    if echo "$OUTPUT" | grep -q "not yet ready"; then
        return 1  # Not ready
    else
        return 0  # Ready
    fi
}

# Main monitoring loop
ELAPSED=0
SSH_READY=0

echo "Starting pod monitoring..."
echo ""

while [ $ELAPSED -lt $MAX_WAIT ]; do
    TIMESTAMP=$(date '+%Y-%m-%d %H:%M:%S')

    # Check pod status
    echo "[$TIMESTAMP] Checking pod status..."
    POD_STATUS=$(check_pod_status)

    if echo "$POD_STATUS" | grep -q "RUNNING"; then
        echo -e "${GREEN}✓ Pod is RUNNING${NC}"

        # Check SSH readiness
        if check_ssh_ready; then
            echo -e "${GREEN}✓ SSH is READY${NC}"
            SSH_READY=1

            # Get SSH connection command
            echo ""
            echo "========================================="
            echo "SSH CONNECTION AVAILABLE"
            echo "========================================="
            runpodctl ssh connect "$POD_ID"
            echo "========================================="
            echo ""
            echo "Next steps:"
            echo "1. Use the SSH command above to connect to the pod"
            echo "2. Run validation commands:"
            echo "   - ls -lh /runpod-volume/"
            echo "   - nvidia-smi"
            echo "   - ps aux | grep train_tft_parquet"
            echo "   - tail -f /var/log/entrypoint.log (if logged)"
            echo ""
            echo "3. Or use web UI for logs:"
            echo "   https://www.runpod.io/console/pods"
            echo ""

            break  # Exit loop
        else
            echo -e "${YELLOW}⏳ SSH not yet ready (initializing...)${NC}"
        fi

    elif echo "$POD_STATUS" | grep -q "STOPPED"; then
        echo -e "${RED}✗ Pod has STOPPED${NC}"
        echo ""
        echo "Pod may have completed or crashed."
        echo "Check logs at: https://www.runpod.io/console/pods"
        echo ""
        exit 1

    else
        echo -e "${YELLOW}⏳ Pod status: UNKNOWN${NC}"
        echo "$POD_STATUS"
    fi

    # Show elapsed time and wait
    ELAPSED=$((ELAPSED + CHECK_INTERVAL))
    REMAINING=$((MAX_WAIT - ELAPSED))
    echo "Elapsed: ${ELAPSED}s | Remaining: ${REMAINING}s"
    echo ""

    if [ $ELAPSED -lt $MAX_WAIT ]; then
        echo "Waiting ${CHECK_INTERVAL}s before next check..."
        sleep $CHECK_INTERVAL
        echo ""
    fi
done

# Final status
echo "========================================="
echo "MONITORING SUMMARY"
echo "========================================="
echo "Pod ID: $POD_ID"
echo "Total Time: ${ELAPSED}s ($(($ELAPSED / 60)) minutes)"

if [ $SSH_READY -eq 1 ]; then
    echo -e "Status: ${GREEN}SUCCESS - SSH Ready${NC}"
    echo ""
    echo "Pod is ready for validation!"
    echo "Use the SSH connection command above or web UI logs."
else
    echo -e "Status: ${YELLOW}TIMEOUT - SSH Not Ready${NC}"
    echo ""
    echo "SSH service did not become ready within $((MAX_WAIT / 60)) minutes."
    echo ""
    echo "Possible causes:"
    echo "1. Pod is still initializing (wait longer)"
    echo "2. SSH service disabled in Docker image"
    echo "3. Pod crashed during startup"
    echo ""
    echo "Recommended actions:"
    echo "1. Check web UI logs: https://www.runpod.io/console/pods"
    echo "2. Verify pod status: runpodctl get pod $POD_ID"
    echo "3. Check if training completed (may have run before SSH ready)"
fi

echo "========================================="
echo ""
echo "Full pod details:"
runpodctl get pod "$POD_ID"
echo ""
echo "To terminate pod when done:"
echo "  runpodctl remove pod $POD_ID"
echo ""