#!/bin/bash
set -e

# =============================================================================
# RUNPOD S3 UPLOAD SCRIPT
# =============================================================================
# Uploads training binaries and data to Runpod Network Volume via S3 API
# This is a one-time operation - binaries persist until manually deleted
#
# Prerequisites:
#   1. Create Runpod Network Volume (get volume ID)
#   2. Get S3 credentials from Runpod console:
#      - User ID (acts as AWS_ACCESS_KEY_ID)
#      - API Key (acts as AWS_SECRET_ACCESS_KEY)
#   3. Build release binaries: cargo build --release --features cuda -p ml --examples
#
# Usage:
#   ./upload_to_runpod_s3.sh
#
# Or with custom config:
#   S3_ENDPOINT=https://s3api-eu-ro-1.runpod.io \
#   S3_BUCKET=your-volume-id \
#   ./upload_to_runpod_s3.sh
# =============================================================================

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

echo "=========================================="
echo "Runpod S3 Upload Script"
echo "=========================================="

# Load configuration from .env.runpod if exists
if [ -f ".env.runpod" ]; then
    echo "Loading configuration from .env.runpod..."
    source .env.runpod
else
    echo -e "${YELLOW}Warning: .env.runpod not found${NC}"
    echo "Create .env.runpod with:"
    echo "  S3_ENDPOINT=https://s3api-DATACENTER.runpod.io"
    echo "  S3_BUCKET=your-network-volume-id"
    echo "  AWS_ACCESS_KEY_ID=your-runpod-user-id"
    echo "  AWS_SECRET_ACCESS_KEY=your-runpod-api-key"
    echo ""
fi

# Validate required environment variables
required_vars=("S3_ENDPOINT" "S3_BUCKET" "AWS_ACCESS_KEY_ID" "AWS_SECRET_ACCESS_KEY")
missing_vars=()

for var in "${required_vars[@]}"; do
    if [ -z "${!var}" ]; then
        missing_vars+=("$var")
    fi
done

if [ ${#missing_vars[@]} -gt 0 ]; then
    echo -e "${RED}ERROR: Missing required environment variables:${NC}"
    for var in "${missing_vars[@]}"; do
        echo "  - $var"
    done
    echo ""
    echo "Set them via environment or create .env.runpod file"
    exit 1
fi

# Export for AWS CLI
export AWS_REGION="${AWS_REGION:-us-east-1}"
export AWS_ACCESS_KEY_ID
export AWS_SECRET_ACCESS_KEY

echo "Configuration:"
echo "  S3 Endpoint: $S3_ENDPOINT"
echo "  S3 Bucket: $S3_BUCKET"
echo "  AWS Region: $AWS_REGION"
echo "  User ID: ${AWS_ACCESS_KEY_ID:0:10}... (masked)"
echo ""

# Configure AWS CLI for Runpod S3
aws configure set default.s3.signature_version s3v4
aws configure set default.s3.addressing_style path

# Test S3 connection
echo "Testing S3 connection..."
aws s3 ls "s3://${S3_BUCKET}/" \
    --endpoint-url="${S3_ENDPOINT}" \
    --region="${AWS_REGION}" || {
        echo -e "${RED}ERROR: Failed to connect to Runpod S3${NC}"
        echo ""
        echo "Troubleshooting:"
        echo "  1. Verify S3_ENDPOINT matches your datacenter (e.g., us-ca-1, eu-ro-1)"
        echo "  2. Verify S3_BUCKET is your Network Volume ID (not a custom name)"
        echo "  3. Verify AWS_ACCESS_KEY_ID is your Runpod User ID"
        echo "  4. Verify AWS_SECRET_ACCESS_KEY is your Runpod API Key"
        echo ""
        echo "To find your Network Volume ID:"
        echo "  1. Go to https://www.runpod.io/console/storage"
        echo "  2. Click on your Network Volume"
        echo "  3. Copy the Volume ID (looks like: abc123xyz456)"
        exit 1
    }

echo -e "${GREEN}✓ S3 connection successful${NC}"
echo ""

# Define binaries and data files to upload
BINARIES=(
    "target/release/examples/train_tft_parquet"
    "target/release/examples/train_mamba2_parquet"
    "target/release/examples/train_dqn"
    "target/release/examples/train_ppo"
)

DATA_FILES=(
    "test_data/ES_FUT_180d.parquet"
)

# Upload binaries
echo "=========================================="
echo "Uploading Training Binaries"
echo "=========================================="

for binary in "${BINARIES[@]}"; do
    binary_name=$(basename "$binary")

    if [ ! -f "$binary" ]; then
        echo -e "${YELLOW}Warning: Binary not found: $binary${NC}"
        echo "         Run: cargo build --release --features cuda -p ml --examples"
        continue
    fi

    binary_size=$(stat -c%s "$binary" 2>/dev/null || stat -f%z "$binary")
    binary_size_mb=$(echo "scale=2; $binary_size / 1024 / 1024" | bc)

    echo ""
    echo "Uploading: $binary_name (${binary_size_mb} MB)"

    aws s3 cp \
        "$binary" \
        "s3://${S3_BUCKET}/binaries/${binary_name}" \
        --endpoint-url="${S3_ENDPOINT}" \
        --region="${AWS_REGION}" || {
            echo -e "${RED}ERROR: Failed to upload $binary_name${NC}"
            exit 1
        }

    echo -e "${GREEN}✓ Uploaded successfully${NC}"
done

# Upload data files
echo ""
echo "=========================================="
echo "Uploading Training Data"
echo "=========================================="

for data_file in "${DATA_FILES[@]}"; do
    data_name=$(basename "$data_file")

    if [ ! -f "$data_file" ]; then
        echo -e "${YELLOW}Warning: Data file not found: $data_file${NC}"
        echo "         This is optional - you can mount data as volume instead"
        continue
    fi

    data_size=$(stat -c%s "$data_file" 2>/dev/null || stat -f%z "$data_file")
    data_size_mb=$(echo "scale=2; $data_size / 1024 / 1024" | bc)

    echo ""
    echo "Uploading: $data_name (${data_size_mb} MB)"

    aws s3 cp \
        "$data_file" \
        "s3://${S3_BUCKET}/data/${data_name}" \
        --endpoint-url="${S3_ENDPOINT}" \
        --region="${AWS_REGION}" || {
            echo -e "${RED}ERROR: Failed to upload $data_name${NC}"
            exit 1
        }

    echo -e "${GREEN}✓ Uploaded successfully${NC}"
done

# List uploaded files
echo ""
echo "=========================================="
echo "Uploaded Files"
echo "=========================================="

echo ""
echo "Binaries:"
aws s3 ls "s3://${S3_BUCKET}/binaries/" \
    --endpoint-url="${S3_ENDPOINT}" \
    --region="${AWS_REGION}" || true

echo ""
echo "Data Files:"
aws s3 ls "s3://${S3_BUCKET}/data/" \
    --endpoint-url="${S3_ENDPOINT}" \
    --region="${AWS_REGION}" || true

echo ""
echo "=========================================="
echo -e "${GREEN}Upload Complete!${NC}"
echo "=========================================="
echo ""
echo "Next steps:"
echo "  1. Build Docker image: docker build -f Dockerfile.runpod.s3 -t foxhunt-runpod-s3:latest ."
echo "  2. Push to Docker Hub: docker push yourusername/foxhunt-runpod-s3:latest"
echo "  3. Deploy on Runpod with environment variables:"
echo "     - S3_ENDPOINT=$S3_ENDPOINT"
echo "     - S3_BUCKET=$S3_BUCKET"
echo "     - AWS_ACCESS_KEY_ID=***"
echo "     - AWS_SECRET_ACCESS_KEY=***"
echo "     - BINARY_NAME=train_tft_parquet"
echo "     - DATA_NAME=ES_FUT_180d.parquet"
echo ""