#!/bin/bash
################################################################################
# runpod_upload.sh - Upload Foxhunt binaries and data to Runpod Network Volume
#
# Purpose:
#   Builds release binaries with CUDA support and uploads them to Runpod
#   Network Volume using AWS S3-compatible API. This script runs on LOCAL
#   CLIENT ONLY. Runpod pods will mount the volume and access files at
#   /runpod-volume/.
#
# Prerequisites:
#   1. AWS CLI installed: apt-get install awscli
#   2. Runpod credentials configured in ~/.aws/credentials:
#      [runpod]
#      aws_access_key_id = <your-runpod-user-id>
#      aws_secret_access_key = <your-runpod-api-key>
#   3. Environment variable RUNPOD_S3_ENDPOINT set to your Runpod endpoint
#      Example: export RUNPOD_S3_ENDPOINT="https://s3api-eur-is-1.runpod.io"
#
# Environment Variables:
#   - RUNPOD_S3_ENDPOINT: Runpod S3-compatible endpoint (required)
#   - AWS_PROFILE: AWS profile to use (default: runpod)
#
# Usage:
#   export RUNPOD_S3_ENDPOINT="https://s3api-eur-is-1.runpod.io"
#   ./scripts/runpod_upload.sh
#
# What gets uploaded:
#   - Release binaries: train_tft_parquet, train_dqn, train_ppo, train_mamba2_dbn
#   - Test data: 9 Parquet files (ES.FUT, NQ.FUT, 6E.FUT, ZN.FUT variants)
#   - Directory structure: /runpod-volume/{bin,test_data,models}
#
# Post-upload access on Runpod pod:
#   - Binaries: /runpod-volume/bin/train_tft_parquet
#   - Data: /runpod-volume/test_data/ES_FUT_180d.parquet
#   - Models: /runpod-volume/models/ (for trained output)
################################################################################

set -e  # Exit on any error

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

################################################################################
# Configuration Validation
################################################################################

echo -e "${BOLD}=========================================${NC}"
echo -e "${BOLD}Foxhunt Runpod Upload Script${NC}"
echo -e "${BOLD}=========================================${NC}"
echo ""

# Validate RUNPOD_S3_ENDPOINT environment variable
if [ -z "$RUNPOD_S3_ENDPOINT" ]; then
    echo -e "${RED}ERROR: RUNPOD_S3_ENDPOINT environment variable not set${NC}"
    echo ""
    echo "Please set your Runpod S3 endpoint:"
    echo "  export RUNPOD_S3_ENDPOINT=\"https://s3api-<datacenter>.runpod.io\""
    echo ""
    echo "Find your endpoint in Runpod console:"
    echo "  1. Go to https://www.runpod.io/console/user/settings"
    echo "  2. Navigate to 'Network Volumes'"
    echo "  3. Copy the S3 endpoint URL"
    exit 1
fi

# Set AWS profile (default: runpod)
AWS_PROFILE="${AWS_PROFILE:-runpod}"

# Validate AWS profile exists
if ! aws configure list --profile "$AWS_PROFILE" &>/dev/null; then
    echo -e "${RED}ERROR: AWS profile '$AWS_PROFILE' not found${NC}"
    echo ""
    echo "Please configure Runpod credentials in ~/.aws/credentials:"
    echo "  [runpod]"
    echo "  aws_access_key_id = <your-runpod-user-id>"
    echo "  aws_secret_access_key = <your-runpod-api-key>"
    echo ""
    echo "Find your credentials in Runpod console:"
    echo "  1. Go to https://www.runpod.io/console/user/settings"
    echo "  2. Navigate to 'API Keys'"
    echo "  3. Copy User ID and API Key"
    exit 1
fi

# Configuration summary
echo -e "${BLUE}Configuration:${NC}"
echo "  Endpoint: $RUNPOD_S3_ENDPOINT"
echo "  AWS Profile: $AWS_PROFILE"
echo "  Upload Target: s3://runpod-volume/foxhunt/"
echo ""

# Project root directory
FOXHUNT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$FOXHUNT_ROOT"

################################################################################
# Step 1: Build Release Binaries with CUDA
################################################################################

echo -e "${BOLD}Step 1/5: Building release binaries with CUDA support...${NC}"
echo "This may take 5-10 minutes depending on system."
echo ""

START_TIME=$(date +%s)

# Build with release profile and CUDA features
if cargo build --release --features cuda --workspace; then
    END_TIME=$(date +%s)
    BUILD_DURATION=$((END_TIME - START_TIME))
    BUILD_MIN=$(awk "BEGIN {printf \"%.2f\", $BUILD_DURATION/60}")
    echo -e "${GREEN}✅ Build complete in ${BUILD_MIN} minutes${NC}"
else
    echo -e "${RED}❌ ERROR: Cargo build failed${NC}"
    echo "Debug steps:"
    echo "  1. Check Rust toolchain: rustc --version"
    echo "  2. Verify CUDA installation: nvidia-smi"
    echo "  3. Check build logs above for specific errors"
    exit 1
fi

echo ""

################################################################################
# Step 2: Upload Training Binaries
################################################################################

echo -e "${BOLD}Step 2/5: Uploading training binaries...${NC}"
echo ""

# List of training binaries to upload
BINARIES=(
    "train_tft_parquet"
    "train_dqn"
    "train_ppo"
    "train_mamba2_dbn"
)

TOTAL_BIN_SIZE=0
UPLOADED_COUNT=0

for binary in "${BINARIES[@]}"; do
    SRC="$FOXHUNT_ROOT/target/release/examples/$binary"

    if [ -f "$SRC" ]; then
        # Get file size
        SIZE=$(stat -c%s "$SRC" 2>/dev/null || stat -f%z "$SRC" 2>/dev/null)
        SIZE_MB=$(awk "BEGIN {printf \"%.2f\", $SIZE/1024/1024}")
        TOTAL_BIN_SIZE=$((TOTAL_BIN_SIZE + SIZE))

        # Upload to Runpod S3
        DEST="s3://se3zdnb5o4/binaries/$binary"

        if aws s3 cp "$SRC" "$DEST" \
            --profile "$AWS_PROFILE" \
            --endpoint-url "$RUNPOD_S3_ENDPOINT"; then
            echo -e "  ${GREEN}✓${NC} $binary (${SIZE_MB} MB)"
            UPLOADED_COUNT=$((UPLOADED_COUNT + 1))
        else
            echo -e "  ${RED}✗${NC} $binary upload failed"
        fi
    else
        echo -e "  ${YELLOW}⚠${NC} $binary not found (skipping)"
    fi
done

TOTAL_BIN_MB=$(awk "BEGIN {printf \"%.2f\", $TOTAL_BIN_SIZE/1024/1024}")
echo ""
echo "Uploaded $UPLOADED_COUNT binaries (${TOTAL_BIN_MB} MB total)"
echo ""

################################################################################
# Step 3: Upload Test Data (Parquet Files)
################################################################################

echo -e "${BOLD}Step 3/5: Uploading test data (Parquet files)...${NC}"
echo ""

PARQUET_FILES=(
    "ES_FUT_180d.parquet"
    "NQ_FUT_180d.parquet"
    "6E_FUT_180d.parquet"
    "ZN_FUT_90d.parquet"
    "ZN_FUT_90d_clean.parquet"
    "ES_FUT_small.parquet"
    "NQ_FUT_small.parquet"
    "6E_FUT_small.parquet"
    "ZN_FUT_small.parquet"
)

TOTAL_DATA_SIZE=0
UPLOADED_DATA_COUNT=0

for file in "${PARQUET_FILES[@]}"; do
    SRC="$FOXHUNT_ROOT/test_data/$file"

    if [ -f "$SRC" ]; then
        # Get file size
        SIZE=$(stat -c%s "$SRC" 2>/dev/null || stat -f%z "$SRC" 2>/dev/null)
        SIZE_MB=$(awk "BEGIN {printf \"%.2f\", $SIZE/1024/1024}")
        TOTAL_DATA_SIZE=$((TOTAL_DATA_SIZE + SIZE))

        # Upload to Runpod S3
        DEST="s3://se3zdnb5o4/test_data/$file"

        if aws s3 cp "$SRC" "$DEST" \
            --profile "$AWS_PROFILE" \
            --endpoint-url "$RUNPOD_S3_ENDPOINT"; then
            echo -e "  ${GREEN}✓${NC} $file (${SIZE_MB} MB)"
            UPLOADED_DATA_COUNT=$((UPLOADED_DATA_COUNT + 1))
        else
            echo -e "  ${RED}✗${NC} $file upload failed"
        fi
    else
        echo -e "  ${YELLOW}⚠${NC} $file not found (skipping)"
    fi
done

TOTAL_DATA_MB=$(awk "BEGIN {printf \"%.2f\", $TOTAL_DATA_SIZE/1024/1024}")
echo ""
echo "Uploaded $UPLOADED_DATA_COUNT data files (${TOTAL_DATA_MB} MB total)"
echo ""

################################################################################
# Step 4: Create Models Directory
################################################################################

echo -e "${BOLD}Step 4/5: Creating /runpod-volume/models/ directory...${NC}"
echo ""

# Create empty marker file to ensure directory exists
MARKER_FILE=$(mktemp)
echo "Model storage directory created by runpod_upload.sh" > "$MARKER_FILE"
echo "Date: $(date)" >> "$MARKER_FILE"

if aws s3 cp "$MARKER_FILE" "s3://se3zdnb5o4/models/.directory_created" \
    --profile "$AWS_PROFILE" \
    --endpoint-url "$RUNPOD_S3_ENDPOINT"; then
    echo -e "${GREEN}✅ Models directory created${NC}"
else
    echo -e "${YELLOW}⚠ Warning: Could not create models directory marker${NC}"
fi

rm -f "$MARKER_FILE"
echo ""

################################################################################
# Step 5: Verify Uploads
################################################################################

echo -e "${BOLD}Step 5/5: Verifying uploads...${NC}"
echo ""

# List uploaded files to verify
echo "Verifying binaries/ directory:"
if aws s3 ls "s3://se3zdnb5o4/binaries/" \
    --profile "$AWS_PROFILE" \
    --endpoint-url "$RUNPOD_S3_ENDPOINT" | head -10; then
    echo -e "${GREEN}✅ Binaries verified${NC}"
else
    echo -e "${RED}❌ Could not verify binaries${NC}"
fi

echo ""
echo "Verifying test_data/ directory:"
if aws s3 ls "s3://se3zdnb5o4/test_data/" \
    --profile "$AWS_PROFILE" \
    --endpoint-url "$RUNPOD_S3_ENDPOINT" | head -10; then
    echo -e "${GREEN}✅ Test data verified${NC}"
else
    echo -e "${RED}❌ Could not verify test data${NC}"
fi

echo ""
echo "Verifying models/ directory:"
if aws s3 ls "s3://se3zdnb5o4/models/" \
    --profile "$AWS_PROFILE" \
    --endpoint-url "$RUNPOD_S3_ENDPOINT" | head -5; then
    echo -e "${GREEN}✅ Models directory verified${NC}"
else
    echo -e "${YELLOW}⚠ Warning: Models directory not found (non-critical)${NC}"
fi

################################################################################
# Upload Summary
################################################################################

echo ""
echo -e "${BOLD}=========================================${NC}"
echo -e "${BOLD}${GREEN}✅ Upload Complete!${NC}${BOLD}${NC}"
echo -e "${BOLD}=========================================${NC}"
echo ""
echo -e "${BLUE}Upload Summary:${NC}"
echo "  Binaries: $UPLOADED_COUNT files (${TOTAL_BIN_MB} MB)"
echo "  Test Data: $UPLOADED_DATA_COUNT files (${TOTAL_DATA_MB} MB)"
TOTAL_SIZE_MB=$(awk "BEGIN {printf \"%.2f\", ($TOTAL_BIN_SIZE + $TOTAL_DATA_SIZE)/1024/1024}")
echo "  Total: ${TOTAL_SIZE_MB} MB"
echo ""

echo -e "${BLUE}Files available at (on Runpod pod):${NC}"
echo "  Binaries: /runpod-volume/binaries/"
echo "  Test Data: /runpod-volume/test_data/"
echo "  Models (output): /runpod-volume/models/"
echo ""

echo -e "${BLUE}Next Steps:${NC}"
echo "1. Create Runpod pod with Tesla V100 16GB GPU"
echo "2. Mount network volume (se3zdnb5o4) to /runpod-volume"
echo "3. Run training on pod:"
echo "   cd /runpod-volume"
echo "   ./binaries/train_tft_parquet \\"
echo "     --parquet-file ./test_data/ES_FUT_180d.parquet \\"
echo "     --epochs 50 \\"
echo "     --output-dir ./models/tft_fp32"
echo ""

echo -e "${BLUE}To download trained models:${NC}"
echo "  aws s3 sync \\"
echo "    --profile runpod \\"
echo "    --endpoint-url $RUNPOD_S3_ENDPOINT \\"
echo "    s3://se3zdnb5o4/models/ ./models/"
echo ""

echo -e "${BOLD}=========================================${NC}"
echo -e "${GREEN}Ready for Runpod deployment!${NC}"
echo -e "${BOLD}=========================================${NC}"