#!/bin/bash
# =============================================================================
# FOXHUNT FP32 PRODUCTION DEPLOYMENT SCRIPT
# =============================================================================
# Purpose: Build all ML training binaries with full optimizations and CUDA support
# Target: Runpod GPU deployment with volume mount architecture
# Build Time: ~6 minutes (with mimalloc allocator)
# Output: 4 production-ready binaries (~160MB total)
# =============================================================================

set -euo pipefail

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

# Timestamp function
timestamp() {
    date '+%Y-%m-%d %H:%M:%S'
}

# Header
echo -e "${BLUE}=============================================${NC}"
echo -e "${BLUE}Foxhunt FP32 Production Deployment${NC}"
echo -e "${BLUE}=============================================${NC}"
echo "[$(timestamp)] Starting production build process..."
echo ""

# =============================================================================
# PRE-FLIGHT CHECKS
# =============================================================================

echo -e "${YELLOW}[1/5] Pre-flight checks${NC}"
echo "[$(timestamp)] Verifying build environment..."

# Check CUDA installation
if command -v nvcc &> /dev/null; then
    CUDA_VERSION=$(nvcc --version | grep "release" | awk '{print $5}' | cut -d',' -f1)
    echo "  ✓ CUDA compiler found: ${CUDA_VERSION}"
else
    echo -e "  ${RED}✗ CUDA compiler (nvcc) not found${NC}"
    echo "    Install CUDA Toolkit 12.0+ for GPU acceleration"
    exit 1
fi

# Check cargo
if ! command -v cargo &> /dev/null; then
    echo -e "  ${RED}✗ Cargo not found${NC}"
    echo "    Install Rust: https://rustup.rs/"
    exit 1
fi
echo "  ✓ Cargo found: $(cargo --version)"

# Check workspace
if [ ! -f "Cargo.toml" ]; then
    echo -e "  ${RED}✗ Not in Foxhunt workspace root${NC}"
    echo "    Run this script from the project root directory"
    exit 1
fi
echo "  ✓ Workspace root verified"

# Verify ml crate exists
if [ ! -d "ml" ]; then
    echo -e "  ${RED}✗ ML crate not found${NC}"
    exit 1
fi
echo "  ✓ ML crate found"

echo ""

# =============================================================================
# BUILD ALL ML TRAINING BINARIES
# =============================================================================

echo -e "${YELLOW}[2/5] Building ML training binaries${NC}"
echo "[$(timestamp)] Building with optimizations: CUDA + mimalloc + release"
echo ""

# Build parameters
CARGO_FEATURES="cuda,mimalloc-allocator"
BUILD_FLAGS="--release"
CRATE="-p ml"
EXAMPLES="--examples"

# Build command
BUILD_CMD="cargo build ${BUILD_FLAGS} ${CRATE} --features ${CARGO_FEATURES} ${EXAMPLES}"

echo "Build command:"
echo "  ${BUILD_CMD}"
echo ""
echo "Optimizations enabled:"
echo "  • CUDA GPU acceleration (12.x+)"
echo "  • mimalloc allocator (faster memory management)"
echo "  • Release mode (full LTO, optimizations)"
echo "  • All ML examples (TFT, MAMBA-2, DQN, PPO)"
echo ""

# Start build with timing
BUILD_START=$(date +%s)

if ${BUILD_CMD}; then
    BUILD_END=$(date +%s)
    BUILD_TIME=$((BUILD_END - BUILD_START))
    BUILD_MINS=$((BUILD_TIME / 60))
    BUILD_SECS=$((BUILD_TIME % 60))
    
    echo ""
    echo -e "${GREEN}✓ Build completed successfully${NC}"
    echo "  Build time: ${BUILD_MINS}m ${BUILD_SECS}s"
else
    echo -e "${RED}✗ Build failed${NC}"
    echo "  Check errors above and fix compilation issues"
    exit 1
fi

echo ""

# =============================================================================
# VERIFY BINARIES
# =============================================================================

echo -e "${YELLOW}[3/5] Verifying binaries${NC}"
echo "[$(timestamp)] Checking output binaries..."
echo ""

# Expected binaries
BINARIES=(
    "train_tft_parquet"
    "train_mamba2_parquet"
    "train_dqn"
    "train_ppo"
)

BINARY_DIR="target/release/examples"
TOTAL_SIZE=0
ALL_FOUND=true

for binary in "${BINARIES[@]}"; do
    BINARY_PATH="${BINARY_DIR}/${binary}"
    
    if [ -f "${BINARY_PATH}" ]; then
        SIZE=$(stat -c%s "${BINARY_PATH}" 2>/dev/null || stat -f%z "${BINARY_PATH}")
        SIZE_MB=$((SIZE / 1024 / 1024))
        TOTAL_SIZE=$((TOTAL_SIZE + SIZE))
        
        # Verify executable
        if [ -x "${BINARY_PATH}" ]; then
            echo "  ✓ ${binary} (${SIZE_MB}MB, executable)"
        else
            echo -e "  ${YELLOW}⚠ ${binary} (${SIZE_MB}MB, NOT executable)${NC}"
            chmod +x "${BINARY_PATH}"
            echo "    Fixed: Made executable"
        fi
    else
        echo -e "  ${RED}✗ ${binary} (NOT FOUND)${NC}"
        ALL_FOUND=false
    fi
done

TOTAL_SIZE_MB=$((TOTAL_SIZE / 1024 / 1024))
echo ""
echo "Total binary size: ${TOTAL_SIZE_MB}MB"

if [ "$ALL_FOUND" = false ]; then
    echo -e "${RED}✗ Some binaries missing - build may have failed${NC}"
    exit 1
fi

echo ""

# =============================================================================
# DISPLAY DEPLOYMENT INSTRUCTIONS
# =============================================================================

echo -e "${YELLOW}[4/5] Deployment instructions${NC}"
echo ""
echo "Binaries are ready for Runpod deployment!"
echo ""
echo -e "${BLUE}--- RUNPOD VOLUME UPLOAD ---${NC}"
echo "Upload binaries to Runpod Network Volume (one-time setup):"
echo ""
echo "  1. SSH into any Runpod pod with your volume mounted:"
echo "     ssh root@\${POD_ID}.ssh.runpod.io"
echo ""
echo "  2. Create binaries directory (if not exists):"
echo "     mkdir -p /runpod-volume/binaries"
echo ""
echo "  3. Upload binaries from your local machine:"
echo "     scp ${BINARY_DIR}/train_* root@\${POD_ID}.ssh.runpod.io:/runpod-volume/binaries/"
echo ""
echo "  4. Verify upload:"
echo "     ssh root@\${POD_ID}.ssh.runpod.io 'ls -lh /runpod-volume/binaries/'"
echo ""
echo -e "${BLUE}--- RUNPOD POD DEPLOYMENT ---${NC}"
echo "Deploy training pod (automated script):"
echo ""
echo "  # Quick smoke test (DQN 1 epoch)"
echo "  ./scripts/runpod_deploy.py --datacenter EUR-IS-1 --dry-run"
echo "  ./scripts/runpod_deploy.py --datacenter EUR-IS-1"
echo ""
echo "  # TFT-225 production training (50 epochs)"
echo "  ./scripts/runpod_deploy.py --datacenter EUR-IS-1 \\"
echo "    --command '/runpod-volume/binaries/train_tft_parquet --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50 --use-gpu'"
echo ""
echo "  # MAMBA-2 training"
echo "  ./scripts/runpod_deploy.py --datacenter EUR-IS-1 \\"
echo "    --command '/runpod-volume/binaries/train_mamba2_parquet --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50 --use-gpu'"
echo ""
echo -e "${BLUE}--- MANUAL DEPLOYMENT ---${NC}"
echo "Manual Runpod console deployment:"
echo "  • GPU: Tesla V100-PCIE-16GB or RTX A4000 (16GB+ VRAM)"
echo "  • Region: EUR-IS-1 (CRITICAL: matches volume location)"
echo "  • Image: jgrusewski/foxhunt:latest"
echo "  • Volume: Mount your Runpod Network Volume at /runpod-volume"
echo "  • Env: BINARY_NAME=train_tft_parquet (or other model)"
echo "  • Args: --parquet-file /runpod-volume/test_data/ES_FUT_180d.parquet --epochs 50 --use-gpu"
echo ""

# =============================================================================
# COST ESTIMATES
# =============================================================================

echo -e "${YELLOW}[5/5] Cost estimates${NC}"
echo ""
echo "Estimated Runpod costs (Tesla V100 @ \$0.29/hr):"
echo ""
echo "  • DQN smoke test (1 epoch):       ~\$0.01 (~2 minutes)"
echo "  • PPO training (100 epochs):      ~\$0.05 (~10 minutes)"
echo "  • MAMBA-2 training (50 epochs):   ~\$0.10 (~20 minutes)"
echo "  • TFT-225 training (50 epochs):   ~\$0.50 (~100 minutes)"
echo ""
echo "  Daily budget (4 training runs):   ~\$0.66"
echo "  Monthly budget (120 runs):        ~\$20"
echo ""
echo "Network Volume cost: \$5/month (50GB)"
echo ""
echo -e "${BLUE}Total estimated monthly cost: \$25-\$30${NC}"
echo ""

# =============================================================================
# SUCCESS SUMMARY
# =============================================================================

echo -e "${GREEN}=============================================${NC}"
echo -e "${GREEN}✓ Deployment preparation complete${NC}"
echo -e "${GREEN}=============================================${NC}"
echo ""
echo "Next steps:"
echo "  1. Upload binaries to Runpod volume (see instructions above)"
echo "  2. Run deployment script: ./scripts/runpod_deploy.py --datacenter EUR-IS-1"
echo "  3. Monitor training via Runpod console logs"
echo "  4. Download trained models from /runpod-volume/models/"
echo ""
echo "Quick reference: See DEPLOYMENT_COMMANDS.md"
echo ""
echo "[$(timestamp)] Script completed successfully"