#!/bin/bash
# Upload trained model checkpoints to S3 (MinIO)
#
# This script uploads all safetensors checkpoint files from the production
# trained_models directory to the S3 bucket with proper organization.

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
CHECKPOINT_DIR="${PROJECT_ROOT}/ml/trained_models/production"
BUCKET="foxhunt-ml-models"

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

echo "╔══════════════════════════════════════════════════════════╗"
echo "║       Checkpoint Upload to S3 (MinIO)                   ║"
echo "╚══════════════════════════════════════════════════════════╝"
echo ""

# Check if MinIO container is running
if ! docker ps | grep -q foxhunt-minio; then
    echo -e "${RED}ERROR: MinIO container 'foxhunt-minio' is not running${NC}"
    echo "Start it with: docker-compose up -d minio"
    exit 1
fi

# Configure MinIO client
echo "Configuring MinIO client..."
docker exec foxhunt-minio mc alias set local http://localhost:9000 foxhunt foxhunt_dev_password > /dev/null 2>&1

# Check if bucket exists, create if not
if ! docker exec foxhunt-minio mc ls local/${BUCKET} > /dev/null 2>&1; then
    echo "Creating bucket: ${BUCKET}"
    docker exec foxhunt-minio mc mb local/${BUCKET}
fi

# Count checkpoint files
TOTAL_FILES=$(find "${CHECKPOINT_DIR}" -name "*.safetensors" -type f | wc -l)
echo -e "${GREEN}Found ${TOTAL_FILES} checkpoint files${NC}"
echo ""

# Upload statistics
UPLOADED=0
FAILED=0
TOTAL_SIZE=0
START_TIME=$(date +%s)

# Function to parse checkpoint filename and determine S3 path
get_s3_path() {
    local filename="$1"
    local model_name=""
    local version=""

    # Parse filename to extract model and version
    if [[ "$filename" =~ ^(dqn|ppo|mamba2|tft)_.*epoch_?([0-9]+) ]]; then
        model_name="${BASH_REMATCH[1]}"
        version="epoch_${BASH_REMATCH[2]}"
    elif [[ "$filename" =~ ^(dqn|ppo|mamba2|tft)_checkpoint_epoch_?([0-9]+) ]]; then
        model_name="${BASH_REMATCH[1]}"
        version="epoch_${BASH_REMATCH[2]}"
    elif [[ "$filename" =~ ^(dqn|ppo|mamba2|tft)_final ]]; then
        model_name="${BASH_REMATCH[1]}"
        version="final"
    else
        model_name="unknown"
        version="v1.0"
    fi

    echo "${model_name}/${version}/checkpoints/${filename}"
}

# Upload each checkpoint
echo "Uploading checkpoints..."
echo ""

while IFS= read -r checkpoint_file; do
    filename=$(basename "$checkpoint_file")
    s3_path=$(get_s3_path "$filename")
    file_size=$(stat -c%s "$checkpoint_file" 2>/dev/null || stat -f%z "$checkpoint_file" 2>/dev/null)

    echo -n "  Uploading: ${filename} ($(numfmt --to=iec-i --suffix=B $file_size)) -> ${s3_path}... "

    # Copy file to container, then upload, then cleanup
    temp_file="/tmp/${filename}"
    if docker cp "$checkpoint_file" "foxhunt-minio:${temp_file}" > /dev/null 2>&1 && \
       docker exec foxhunt-minio mc cp "${temp_file}" "local/${BUCKET}/${s3_path}" > /dev/null 2>&1 && \
       docker exec foxhunt-minio rm "${temp_file}" > /dev/null 2>&1; then
        echo -e "${GREEN}✓${NC}"
        UPLOADED=$((UPLOADED + 1))
        TOTAL_SIZE=$((TOTAL_SIZE + file_size))
    else
        echo -e "${RED}✗${NC}"
        FAILED=$((FAILED + 1))
        # Cleanup on failure
        docker exec foxhunt-minio rm "${temp_file}" > /dev/null 2>&1 || true
    fi
done < <(find "${CHECKPOINT_DIR}" -name "*.safetensors" -type f | sort)

END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))

# Calculate statistics
TOTAL_SIZE_MB=$((TOTAL_SIZE / 1024 / 1024))
if [ $DURATION -gt 0 ]; then
    THROUGHPUT=$(echo "scale=2; $TOTAL_SIZE_MB / $DURATION" | bc)
else
    THROUGHPUT="N/A"
fi

# Print summary
echo ""
echo "╔══════════════════════════════════════════════════════════╗"
echo "║         Upload Summary                                   ║"
echo "╠══════════════════════════════════════════════════════════╣"
printf "║ Total files:     %-6d                                 ║\n" $TOTAL_FILES
printf "║ Uploaded:        %-6d                                 ║\n" $UPLOADED
printf "║ Failed:          %-6d                                 ║\n" $FAILED
printf "║ Total size:      %-6d MB                             ║\n" $TOTAL_SIZE_MB
printf "║ Duration:        %-6d seconds                         ║\n" $DURATION
if [ "$THROUGHPUT" != "N/A" ]; then
    printf "║ Throughput:      %-6s MB/s                          ║\n" $THROUGHPUT
fi
echo "╚══════════════════════════════════════════════════════════╝"
echo ""

# Verify uploads
echo "Verifying uploads..."
OBJECTS_COUNT=$(docker exec foxhunt-minio mc ls -r local/${BUCKET} | wc -l)
echo -e "${GREEN}S3 bucket now contains ${OBJECTS_COUNT} objects${NC}"

# List bucket structure
echo ""
echo "Bucket structure:"
docker exec foxhunt-minio mc ls local/${BUCKET}/ | head -20

if [ $FAILED -gt 0 ]; then
    echo -e "\n${YELLOW}Warning: ${FAILED} files failed to upload${NC}"
    exit 1
fi

echo -e "\n${GREEN}✓ Upload complete!${NC}"