#!/bin/bash
# Foxhunt Hyperopt Docker Build Script
# Builds CUDA-enabled hyperparameter optimization binaries using cargo-chef

set -euo pipefail

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

# Configuration
IMAGE_NAME="${IMAGE_NAME:-jgrusewski/foxhunt-hyperopt}"
IMAGE_TAG="${IMAGE_TAG:-latest}"
GIT_COMMIT=$(git rev-parse --short HEAD 2>/dev/null || echo "unknown")
BUILD_DATE=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
DOCKERFILE="Dockerfile.foxhunt-build"

echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}Foxhunt Hyperopt Docker Build${NC}"
echo -e "${GREEN}========================================${NC}"
echo ""
echo "Image: ${IMAGE_NAME}:${IMAGE_TAG}"
echo "Commit: ${GIT_COMMIT}"
echo "Build Date: ${BUILD_DATE}"
echo ""

# Check if Docker is running
if ! docker info > /dev/null 2>&1; then
    echo -e "${RED}Error: Docker is not running${NC}"
    exit 1
fi

# Check if Dockerfile exists
if [ ! -f "${DOCKERFILE}" ]; then
    echo -e "${RED}Error: ${DOCKERFILE} not found${NC}"
    exit 1
fi

# Enable BuildKit for better caching
export DOCKER_BUILDKIT=1

echo -e "${YELLOW}Building Docker image...${NC}"
echo ""

# Build with cache and build args
docker build \
    --file "${DOCKERFILE}" \
    --tag "${IMAGE_NAME}:${IMAGE_TAG}" \
    --tag "${IMAGE_NAME}:${GIT_COMMIT}" \
    --build-arg GIT_COMMIT="${GIT_COMMIT}" \
    --build-arg BUILD_DATE="${BUILD_DATE}" \
    --build-arg CUDA_VERSION="12.4.1" \
    --build-arg CUDNN_VERSION="9" \
    --progress=plain \
    .

if [ $? -eq 0 ]; then
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}Build successful!${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo "Image: ${IMAGE_NAME}:${IMAGE_TAG}"
    echo "Image: ${IMAGE_NAME}:${GIT_COMMIT}"
    echo ""
    
    # Get image size
    IMAGE_SIZE=$(docker images "${IMAGE_NAME}:${IMAGE_TAG}" --format "{{.Size}}")
    echo "Image size: ${IMAGE_SIZE}"
    echo ""
    
    # List binaries in image
    echo -e "${YELLOW}Verifying binaries...${NC}"
    docker run --rm "${IMAGE_NAME}:${IMAGE_TAG}" ls -lh /usr/local/bin/hyperopt_*
    echo ""
    
    # Verify GLIBC version
    echo -e "${YELLOW}Verifying GLIBC version...${NC}"
    docker run --rm "${IMAGE_NAME}:${IMAGE_TAG}" ldd --version | head -n 1
    echo ""
    
    echo -e "${GREEN}Next steps:${NC}"
    echo "1. Test locally:"
    echo "   docker run --rm --gpus all ${IMAGE_NAME}:${IMAGE_TAG} hyperopt_dqn_demo --help"
    echo ""
    echo "2. Push to registry:"
    echo "   docker push ${IMAGE_NAME}:${IMAGE_TAG}"
    echo "   docker push ${IMAGE_NAME}:${GIT_COMMIT}"
    echo ""
else
    echo ""
    echo -e "${RED}========================================${NC}"
    echo -e "${RED}Build failed!${NC}"
    echo -e "${RED}========================================${NC}"
    exit 1
fi