#!/bin/bash
# Build script for BackgroundFX Pro Docker images

set -e

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

# Configuration
PROJECT_NAME="backgroundfx-pro"
REGISTRY=""  # Set to your registry URL if pushing to registry
VERSION=$(git describe --tags --always --dirty 2>/dev/null || echo "latest")

# Parse arguments
BUILD_TYPE="gpu"
PUSH=false
NO_CACHE=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --cpu)
            BUILD_TYPE="cpu"
            shift
            ;;
        --all)
            BUILD_TYPE="all"
            shift
            ;;
        --push)
            PUSH=true
            shift
            ;;
        --no-cache)
            NO_CACHE=true
            shift
            ;;
        --registry)
            REGISTRY="$2/"
            shift 2
            ;;
        --version)
            VERSION="$2"
            shift 2
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            echo "Usage: $0 [--cpu|--all] [--push] [--no-cache] [--registry REGISTRY] [--version VERSION]"
            exit 1
            ;;
    esac
done

# Build function
build_image() {
    local dockerfile=$1
    local tag=$2
    local context=${3:-".."}
    
    echo -e "${GREEN}Building ${tag}...${NC}"
    
    CACHE_FLAG=""
    if [ "$NO_CACHE" = true ]; then
        CACHE_FLAG="--no-cache"
    fi
    
    docker build \
        $CACHE_FLAG \
        -f "$dockerfile" \
        -t "${PROJECT_NAME}:${tag}" \
        -t "${PROJECT_NAME}:${tag}-${VERSION}" \
        -t "${REGISTRY}${PROJECT_NAME}:${tag}" \
        -t "${REGISTRY}${PROJECT_NAME}:${tag}-${VERSION}" \
        "$context"
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ Successfully built ${tag}${NC}"
    else
        echo -e "${RED}✗ Failed to build ${tag}${NC}"
        exit 1
    fi
}

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

# Check for NVIDIA Docker runtime (for GPU builds)
if [ "$BUILD_TYPE" != "cpu" ]; then
    if ! docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu20.04 nvidia-smi > /dev/null 2>&1; then
        echo -e "${YELLOW}Warning: NVIDIA Docker runtime not available. GPU features will not work.${NC}"
        echo "Install nvidia-docker2 for GPU support."
    fi
fi

# Build images based on type
echo -e "${GREEN}Starting build process...${NC}"
echo "Version: ${VERSION}"
echo "Build type: ${BUILD_TYPE}"

case $BUILD_TYPE in
    gpu)
        build_image "docker/Dockerfile" "gpu"
        build_image "docker/Dockerfile.prod" "gpu-prod"
        ;;
    cpu)
        build_image "docker/Dockerfile.cpu" "cpu"
        ;;
    all)
        build_image "docker/Dockerfile" "gpu"
        build_image "docker/Dockerfile.prod" "gpu-prod"
        build_image "docker/Dockerfile.cpu" "cpu"
        build_image "docker/Dockerfile.models" "models"
        ;;
esac

# Push to registry if requested
if [ "$PUSH" = true ]; then
    if [ -z "$REGISTRY" ]; then
        echo -e "${YELLOW}Warning: No registry specified, skipping push${NC}"
    else
        echo -e "${GREEN}Pushing images to ${REGISTRY}...${NC}"
        
        case $BUILD_TYPE in
            gpu)
                docker push "${REGISTRY}${PROJECT_NAME}:gpu"
                docker push "${REGISTRY}${PROJECT_NAME}:gpu-${VERSION}"
                docker push "${REGISTRY}${PROJECT_NAME}:gpu-prod"
                docker push "${REGISTRY}${PROJECT_NAME}:gpu-prod-${VERSION}"
                ;;
            cpu)
                docker push "${REGISTRY}${PROJECT_NAME}:cpu"
                docker push "${REGISTRY}${PROJECT_NAME}:cpu-${VERSION}"
                ;;
            all)
                docker push "${REGISTRY}${PROJECT_NAME}:gpu"
                docker push "${REGISTRY}${PROJECT_NAME}:gpu-${VERSION}"
                docker push "${REGISTRY}${PROJECT_NAME}:gpu-prod"
                docker push "${REGISTRY}${PROJECT_NAME}:gpu-prod-${VERSION}"
                docker push "${REGISTRY}${PROJECT_NAME}:cpu"
                docker push "${REGISTRY}${PROJECT_NAME}:cpu-${VERSION}"
                docker push "${REGISTRY}${PROJECT_NAME}:models"
                docker push "${REGISTRY}${PROJECT_NAME}:models-${VERSION}"
                ;;
        esac
        
        echo -e "${GREEN}✓ Images pushed successfully${NC}"
    fi
fi

# Show built images
echo -e "\n${GREEN}Built images:${NC}"
docker images | grep "${PROJECT_NAME}" | head -10

echo -e "\n${GREEN}Build complete!${NC}"
echo "Run with: docker-compose up -d"