#!/bin/bash
# Quick Monitoring Script for MAMBA-2 Hyperopt Deployment
# Pod ID: qlql87w5avv1q1

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="$SCRIPT_DIR/.env.runpod"

# Load environment
if [ -f "$ENV_FILE" ]; then
    export $(grep -v '^#' "$ENV_FILE" | xargs)
else
    echo "ERROR: .env.runpod not found"
    exit 1
fi

POD_ID="qlql87w5avv1q1"

# Function to get pod status
check_status() {
    echo "===================="
    echo "POD STATUS CHECK"
    echo "===================="

    curl -s -X GET \
        "https://rest.runpod.io/v1/pods/$POD_ID" \
        -H "Authorization: Bearer $RUNPOD_API_KEY" \
        | python3 -c "
import sys, json
data = json.load(sys.stdin)
print(f\"Status: {data.get('desiredStatus', 'UNKNOWN')}\")
print(f\"Cost: \${data.get('costPerHr', 0)}/hr\")
print(f\"GPU: {data.get('machine', {}).get('gpuDisplayName', 'TBD')}\")
print(f\"Created: {data.get('createdAt', 'N/A')}\")
runtime = data.get('runtime', {})
if runtime:
    print(f\"Uptime: {runtime.get('uptimeInSeconds', 0)}s\")
    print(f\"Ports: {runtime.get('ports', 'N/A')}\")
else:
    print('Runtime: Pod still provisioning...')
"
    echo ""
}

# Function to try SSH
try_ssh() {
    echo "===================="
    echo "SSH CONNECTION TEST"
    echo "===================="

    if ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o ConnectTimeout=5 \
           -p 19735 root@157.157.221.29 "echo 'Connected!'" 2>/dev/null; then
        echo "✅ SSH is available!"
        echo ""
        echo "Connect with:"
        echo "  ssh -p 19735 root@157.157.221.29"
        echo ""
        return 0
    else
        echo "⏳ SSH not yet available (pod still initializing)"
        echo ""
        return 1
    fi
}

# Function to show training logs (if SSH available)
show_logs() {
    echo "===================="
    echo "TRAINING LOGS (last 50 lines)"
    echo "===================="

    ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        -p 19735 root@157.157.221.29 \
        "tail -50 /workspace/logs/*.log 2>/dev/null || echo 'No logs yet'" 2>/dev/null || \
        echo "Cannot access logs (pod not ready)"
    echo ""
}

# Function to check GPU
check_gpu() {
    echo "===================="
    echo "GPU UTILIZATION"
    echo "===================="

    ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        -p 19735 root@157.157.221.29 \
        "nvidia-smi --query-gpu=name,memory.used,memory.total,utilization.gpu,temperature.gpu --format=csv,noheader" 2>/dev/null || \
        echo "Cannot access GPU info (pod not ready)"
    echo ""
}

# Function to check for critical success metric
check_losses() {
    echo "===================="
    echo "LOSS VALIDATION (CRITICAL)"
    echo "===================="

    ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        -p 19735 root@157.157.221.29 \
        "grep -E 'Feature normalization|Epoch.*Loss' /workspace/logs/*.log 2>/dev/null | tail -20 || echo 'No training logs yet'" 2>/dev/null || \
        echo "Cannot access logs (pod not ready)"
    echo ""
    echo "SUCCESS CRITERIA:"
    echo "  ✅ Train Loss < 1.0"
    echo "  ✅ Val Loss < 1.0"
    echo "  ❌ If losses > 1.0, FIX FAILED"
    echo ""
}

# Main monitoring loop
main() {
    while true; do
        clear
        echo "========================================"
        echo "MAMBA-2 HYPEROPT MONITORING"
        echo "Pod ID: $POD_ID"
        echo "Time: $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
        echo "========================================"
        echo ""

        check_status

        if try_ssh; then
            check_gpu
            check_losses
            show_logs

            echo "===================="
            echo "NEXT ACTIONS"
            echo "===================="
            echo "1. Watch for 'Feature normalization' log line (confirms fix)"
            echo "2. Verify losses < 1.0 (CRITICAL)"
            echo "3. Monitor GPU utilization > 85%"
            echo "4. Check epoch time ~10 min"
            echo ""
            echo "Press Ctrl+C to exit, or wait 60s for refresh..."
            sleep 60
        else
            echo "Pod is still provisioning. Checking again in 30 seconds..."
            sleep 30
        fi
    done
}

# Handle Ctrl+C gracefully
trap 'echo ""; echo "Monitoring stopped."; exit 0' INT

# Parse command line
case "${1:-}" in
    status)
        check_status
        ;;
    ssh)
        try_ssh && echo "To connect:" && echo "ssh -p 19735 root@157.157.221.29"
        ;;
    logs)
        show_logs
        ;;
    gpu)
        check_gpu
        ;;
    losses)
        check_losses
        ;;
    *)
        main
        ;;
esac