#!/bin/bash
# Runpod Configuration Verification Script
# Checks all prerequisites before executing training

set -euo pipefail

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

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Runpod Configuration Verification${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

CHECKS_PASSED=0
CHECKS_FAILED=0

# Check 1: runpodctl installation
echo -n "Checking runpodctl installation... "
if command -v runpodctl &>/dev/null; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: runpodctl not found"
    echo "  Install: curl -fsSL https://raw.githubusercontent.com/runpod/runpodctl/main/install.sh | bash"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 2: API key configuration
echo -n "Checking Runpod API key... "
if runpodctl config 2>&1 | grep -q "apiKey"; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: API key not configured"
    echo "  Get key: https://www.runpod.io/console/user/settings"
    echo "  Configure: runpodctl config --apiKey YOUR_API_KEY"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 3: Training data
echo -n "Checking training data (ES_FUT_180d.parquet)... "
if [ -f "test_data/ES_FUT_180d.parquet" ]; then
    SIZE=$(du -h test_data/ES_FUT_180d.parquet | cut -f1)
    echo -e "${GREEN}✅ (${SIZE})${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: test_data/ES_FUT_180d.parquet not found"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 4: 225 features configured
echo -n "Checking 225-feature configuration... "
if grep -q "features: \[f64; 225\]" ml/src/features/unified.rs; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: 225 features not configured in ml/src/features/unified.rs"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 5: Training example exists
echo -n "Checking TFT training example... "
if [ -f "ml/examples/train_tft_parquet.rs" ]; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: ml/examples/train_tft_parquet.rs not found"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 6: Cargo workspace valid
echo -n "Checking Cargo workspace... "
if cargo check --workspace --quiet 2>/dev/null; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${YELLOW}⚠️ (Has warnings)${NC}"
    echo "  Warning: Cargo check has warnings (non-blocking)"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
fi

# Check 7: Training script exists
echo -n "Checking training script... "
if [ -f "scripts/train_runpod_225_features.sh" ] && [ -x "scripts/train_runpod_225_features.sh" ]; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: scripts/train_runpod_225_features.sh not found or not executable"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 8: Backtesting script exists
echo -n "Checking backtesting script... "
if [ -f "scripts/backtest_runpod_225.sh" ] && [ -x "scripts/backtest_runpod_225.sh" ]; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${RED}❌${NC}"
    echo "  Error: scripts/backtest_runpod_225.sh not found or not executable"
    CHECKS_FAILED=$((CHECKS_FAILED + 1))
fi

# Check 9: CUDA availability (optional)
echo -n "Checking CUDA availability (optional)... "
if command -v nvidia-smi &>/dev/null; then
    GPU_INFO=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader | head -1)
    echo -e "${GREEN}✅ (${GPU_INFO})${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${YELLOW}⚠️ (Local GPU not required for Runpod)${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
fi

# Check 10: Docker services (optional)
echo -n "Checking Docker services (optional)... "
if docker-compose ps 2>/dev/null | grep -q "postgres"; then
    echo -e "${GREEN}✅${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
else
    echo -e "${YELLOW}⚠️ (Not required for Runpod training)${NC}"
    CHECKS_PASSED=$((CHECKS_PASSED + 1))
fi

# Summary
echo ""
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Verification Summary${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""
echo -e "Checks Passed: ${GREEN}${CHECKS_PASSED}${NC}"
echo -e "Checks Failed: ${RED}${CHECKS_FAILED}${NC}"
echo ""

if [ $CHECKS_FAILED -eq 0 ]; then
    echo -e "${GREEN}✅ All checks passed! Ready for Runpod training.${NC}"
    echo ""
    echo "Next steps:"
    echo "  1. Execute training: ./scripts/train_runpod_225_features.sh"
    echo "  2. Monitor progress: tail -f runpod_training_225.log"
    echo "  3. Run backtesting: ./scripts/backtest_runpod_225.sh"
    exit 0
else
    echo -e "${RED}❌ ${CHECKS_FAILED} check(s) failed. Please fix errors before training.${NC}"
    echo ""
    echo "Common fixes:"
    echo "  • API key: runpodctl config --apiKey YOUR_API_KEY"
    echo "  • Training data: Download ES_FUT_180d.parquet to test_data/"
    echo "  • Scripts: chmod +x scripts/*.sh"
    exit 1
fi