#!/bin/bash
# =============================================================================
# PRODUCTION DOCKER BUILD SCRIPT - AUTOMATIC VERSIONING
# =============================================================================
# Purpose: Build and push Docker images with automatic versioning and validation
# Features:
#   - Automatic version tagging (git commit, timestamp, latest)
#   - BuildKit optimization with cache mounts
#   - Binary validation in built image
#   - Image size reporting
#   - Build time measurement
#   - Multi-platform support (future)
#   - Error handling with exit codes
#
# Usage:
#   ./scripts/build_docker_images.sh                    # Build and push with all defaults
#   ./scripts/build_docker_images.sh --no-push          # Build only, skip push
#   ./scripts/build_docker_images.sh --dry-run          # Show what would be built
#   ./scripts/build_docker_images.sh --dockerfile Dockerfile.runpod  # Specify Dockerfile
#   ./scripts/build_docker_images.sh --platform linux/amd64  # Specific platform
#   ./scripts/build_docker_images.sh --skip-validation  # Skip binary validation
#
# Exit codes:
#   0 - Success
#   1 - Build failed
#   2 - Validation failed
#   3 - Push failed
#   4 - Invalid arguments
# =============================================================================

set -euo pipefail

# =============================================================================
# COLOR DEFINITIONS
# =============================================================================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# =============================================================================
# CONFIGURATION
# =============================================================================
DOCKER_REGISTRY="jgrusewski"
IMAGE_NAME="foxhunt-hyperopt"
DEFAULT_DOCKERFILE="Dockerfile.foxhunt-build"
EXPECTED_BINARIES=(
    "hyperopt_baseline_rl"
    "hyperopt_baseline_supervised"
)

# =============================================================================
# SCRIPT VARIABLES
# =============================================================================
DOCKERFILE="${DEFAULT_DOCKERFILE}"
DO_PUSH=true
DRY_RUN=false
PLATFORM=""
SKIP_VALIDATION=false
BUILD_START_TIME=""
BUILD_END_TIME=""

# =============================================================================
# HELPER FUNCTIONS
# =============================================================================

# Print colored output
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_header() {
    echo -e "\n${CYAN}==============================================================================${NC}"
    echo -e "${CYAN}$1${NC}"
    echo -e "${CYAN}==============================================================================${NC}\n"
}

# Calculate time difference
time_diff() {
    local start=$1
    local end=$2
    local diff=$((end - start))
    echo "${diff}s"
}

# Format bytes to human-readable
format_bytes() {
    local bytes=$1
    if [ "$bytes" -lt 1024 ]; then
        echo "${bytes}B"
    elif [ "$bytes" -lt 1048576 ]; then
        echo "$((bytes / 1024))KB"
    elif [ "$bytes" -lt 1073741824 ]; then
        echo "$((bytes / 1048576))MB"
    else
        echo "$((bytes / 1073741824))GB"
    fi
}

# Check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Show usage
usage() {
    cat << EOF
Usage: $(basename "$0") [OPTIONS]

Build and push Docker images with automatic versioning.

OPTIONS:
    --dockerfile FILE       Dockerfile to build (default: ${DEFAULT_DOCKERFILE})
    --no-push              Build only, skip pushing to registry
    --dry-run              Show what would be built without building
    --platform PLATFORM    Build for specific platform (e.g., linux/amd64)
    --skip-validation      Skip binary validation after build
    -h, --help             Show this help message

EXAMPLES:
    $(basename "$0")                                    # Build and push with all defaults
    $(basename "$0") --no-push                          # Build only, no push
    $(basename "$0") --dry-run                          # Show build plan
    $(basename "$0") --dockerfile Dockerfile.runpod     # Specific Dockerfile
    $(basename "$0") --platform linux/amd64             # Specific platform

EXIT CODES:
    0 - Success
    1 - Build failed
    2 - Validation failed
    3 - Push failed
    4 - Invalid arguments
EOF
}

# =============================================================================
# ARGUMENT PARSING
# =============================================================================
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --dockerfile)
                DOCKERFILE="$2"
                shift 2
                ;;
            --no-push)
                DO_PUSH=false
                shift
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --platform)
                PLATFORM="$2"
                shift 2
                ;;
            --skip-validation)
                SKIP_VALIDATION=true
                shift
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                usage
                exit 4
                ;;
        esac
    done
}

# =============================================================================
# VALIDATION FUNCTIONS
# =============================================================================

# Check prerequisites
check_prerequisites() {
    print_header "CHECKING PREREQUISITES"

    # Check Docker
    if ! command_exists docker; then
        print_error "Docker is not installed"
        exit 4
    fi
    print_success "Docker: $(docker --version)"

    # Check git
    if ! command_exists git; then
        print_error "Git is not installed"
        exit 4
    fi
    print_success "Git: $(git --version)"

    # Check if in git repository
    if ! git rev-parse --git-dir > /dev/null 2>&1; then
        print_error "Not in a git repository"
        exit 4
    fi
    print_success "Git repository detected"

    # Check if Dockerfile exists
    if [ ! -f "$DOCKERFILE" ]; then
        print_error "Dockerfile not found: $DOCKERFILE"
        exit 4
    fi
    print_success "Dockerfile: $DOCKERFILE"

    # Check Docker daemon
    if ! docker info > /dev/null 2>&1; then
        print_error "Docker daemon is not running"
        exit 4
    fi
    print_success "Docker daemon running"

    # Check BuildKit support
    if ! docker buildx version > /dev/null 2>&1; then
        print_warning "BuildKit not available, using legacy builder"
    else
        print_success "BuildKit available: $(docker buildx version | head -n1)"
    fi
}

# Get version tags
get_version_tags() {
    print_header "GENERATING VERSION TAGS"

    # Git commit hash (short)
    GIT_COMMIT=$(git rev-parse --short HEAD 2>/dev/null || echo "unknown")
    print_info "Git commit: $GIT_COMMIT"

    # Timestamp (YYYYMMDD_HHMMSS UTC)
    TIMESTAMP=$(date -u +"%Y%m%d_%H%M%S")
    print_info "Timestamp: $TIMESTAMP"

    # Check if working directory is clean
    if ! git diff-index --quiet HEAD -- 2>/dev/null; then
        print_warning "Working directory has uncommitted changes"
        GIT_COMMIT="${GIT_COMMIT}-dirty"
    fi

    # Generate full image tags
    TAG_LATEST="${DOCKER_REGISTRY}/${IMAGE_NAME}:latest"
    TAG_COMMIT="${DOCKER_REGISTRY}/${IMAGE_NAME}:${GIT_COMMIT}"
    TAG_TIMESTAMP="${DOCKER_REGISTRY}/${IMAGE_NAME}:${TIMESTAMP}"

    print_success "Tags generated:"
    echo "  - ${TAG_LATEST}"
    echo "  - ${TAG_COMMIT}"
    echo "  - ${TAG_TIMESTAMP}"
}

# Validate binaries in built image (EMBEDDED BINARIES ARCHITECTURE)
validate_binaries() {
    local image_tag=$1

    print_header "VALIDATING IMAGE"

    # Check that all expected binaries are embedded in /usr/local/bin/
    print_info "Checking embedded binaries in /usr/local/bin/..."
    local missing_binaries=()

    for binary in "${EXPECTED_BINARIES[@]}"; do
        # Override entrypoint to bypass self-termination wrapper during validation
        if ! docker run --rm --entrypoint="" "$image_tag" test -f "/usr/local/bin/$binary"; then
            print_error "Binary not found: /usr/local/bin/$binary"
            missing_binaries+=("$binary")
        else
            print_success "Binary exists: /usr/local/bin/$binary"
        fi
    done

    if [ ${#missing_binaries[@]} -gt 0 ]; then
        print_error "Missing binaries: ${missing_binaries[*]}"
        return 1
    fi

    # Check CUDA libraries
    print_info "Checking CUDA libraries..."
    if ! docker run --rm --entrypoint="" "$image_tag" test -d /usr/local/cuda; then
        print_error "CUDA directory not found in image"
        return 1
    fi
    print_success "CUDA runtime present: /usr/local/cuda"

    # Check cuDNN libraries
    print_info "Checking cuDNN libraries..."
    if ! docker run --rm --entrypoint="" "$image_tag" sh -c 'ldconfig -p | grep -q libcudnn'; then
        print_warning "cuDNN library not found (may be OK if using base image)"
    else
        print_success "cuDNN library present"
    fi

    # Check GLIBC version (should be 2.35 for Ubuntu 22.04)
    print_info "Checking GLIBC version..."
    local glibc_version
    glibc_version=$(docker run --rm --entrypoint="" "$image_tag" ldd --version | head -n1 | grep -oP '\d+\.\d+' | head -n1)
    print_success "GLIBC version: $glibc_version"

    print_success "Image validation complete"
    return 0
}

# Get image size
get_image_size() {
    local image_tag=$1

    print_header "IMAGE SIZE REPORT"

    # Get image size in bytes
    local size_bytes
    size_bytes=$(docker image inspect "$image_tag" --format='{{.Size}}' 2>/dev/null || echo "0")

    # Get human-readable size
    local size_human
    size_human=$(docker image inspect "$image_tag" --format='{{.Size}}' 2>/dev/null | \
        awk '{ sum=$1; hum[1024**3]="GB"; hum[1024**2]="MB"; hum[1024]="KB";
               for (x=1024**3; x>=1024; x/=1024) {
                   if (sum>=x) { printf "%.2f %s\n", sum/x, hum[x]; break }
               } }')

    if [ -z "$size_human" ]; then
        size_human=$(format_bytes "$size_bytes")
    fi

    print_info "Image size: ${size_human}"

    # Show layer sizes
    print_info "Layer breakdown:"
    docker history --human=true --no-trunc "$image_tag" | head -n 10
}

# =============================================================================
# BUILD FUNCTIONS
# =============================================================================

# Build Docker image
build_image() {
    print_header "BUILDING DOCKER IMAGE"

    BUILD_START_TIME=$(date +%s)

    # Prepare build command
    local build_cmd="docker build"
    local build_args=()

    # Enable BuildKit if available
    if command_exists docker-buildx || [ -n "${DOCKER_BUILDKIT:-}" ]; then
        export DOCKER_BUILDKIT=1
        print_info "BuildKit enabled"
    fi

    # Add build arguments
    build_args+=("--build-arg" "GIT_COMMIT=${GIT_COMMIT}")
    build_args+=("--build-arg" "BUILD_DATE=$(date -u +'%Y-%m-%dT%H:%M:%SZ')")
    build_args+=("--build-arg" "VERSION=${GIT_COMMIT}")

    # Add platform if specified
    if [ -n "$PLATFORM" ]; then
        build_args+=("--platform" "$PLATFORM")
        print_info "Building for platform: $PLATFORM"
    fi

    # Add tags
    build_args+=("-t" "$TAG_LATEST")
    build_args+=("-t" "$TAG_COMMIT")
    build_args+=("-t" "$TAG_TIMESTAMP")

    # Add Dockerfile
    build_args+=("-f" "$DOCKERFILE")

    # Add context (current directory)
    build_args+=(".")

    # Show build command
    print_info "Build command:"
    echo "  DOCKER_BUILDKIT=1 ${build_cmd} ${build_args[*]}"
    echo ""

    # Execute build
    if [ "$DRY_RUN" = true ]; then
        print_warning "DRY RUN: Would execute build command above"
        return 0
    fi

    print_info "Starting build..."
    if ${build_cmd} "${build_args[@]}"; then
        BUILD_END_TIME=$(date +%s)
        local build_time
        build_time=$(time_diff "$BUILD_START_TIME" "$BUILD_END_TIME")
        print_success "Build completed in ${build_time}"
        return 0
    else
        print_error "Build failed"
        return 1
    fi
}

# Push images to registry
push_images() {
    print_header "PUSHING IMAGES TO REGISTRY"

    if [ "$DRY_RUN" = true ]; then
        print_warning "DRY RUN: Would push the following tags:"
        echo "  - ${TAG_LATEST}"
        echo "  - ${TAG_COMMIT}"
        echo "  - ${TAG_TIMESTAMP}"
        return 0
    fi

    if [ "$DO_PUSH" = false ]; then
        print_warning "Skipping push (--no-push flag)"
        return 0
    fi

    # Check Docker Hub login
    if ! docker info 2>/dev/null | grep -q "Username: ${DOCKER_REGISTRY}"; then
        print_warning "Not logged in to Docker Hub as ${DOCKER_REGISTRY}"
        print_info "Run: docker login"
        print_warning "Skipping push"
        return 0
    fi

    # Push all tags
    local tags=("$TAG_LATEST" "$TAG_COMMIT" "$TAG_TIMESTAMP")
    for tag in "${tags[@]}"; do
        print_info "Pushing: $tag"
        if docker push "$tag"; then
            print_success "Pushed: $tag"
        else
            print_error "Failed to push: $tag"
            return 1
        fi
    done

    print_success "All images pushed successfully"
    return 0
}

# =============================================================================
# MAIN EXECUTION
# =============================================================================

main() {
    print_header "FOXHUNT DOCKER BUILD SCRIPT"

    # Parse command line arguments
    parse_args "$@"

    # Check prerequisites
    check_prerequisites

    # Generate version tags
    get_version_tags

    # Show dry-run notice
    if [ "$DRY_RUN" = true ]; then
        print_warning "DRY RUN MODE - No actual changes will be made"
        echo ""
    fi

    # Build image
    if ! build_image; then
        print_error "Build failed"
        exit 1
    fi

    # Skip validation and reporting in dry-run mode
    if [ "$DRY_RUN" = true ]; then
        print_success "Dry-run completed successfully"
        exit 0
    fi

    # Validate binaries (unless skipped)
    if [ "$SKIP_VALIDATION" = false ]; then
        if ! validate_binaries "$TAG_LATEST"; then
            print_error "Validation failed"
            exit 2
        fi
    else
        print_warning "Skipping validation (--skip-validation flag)"
    fi

    # Report image size
    get_image_size "$TAG_LATEST"

    # Push images
    if ! push_images; then
        print_error "Push failed"
        exit 3
    fi

    # Final summary
    print_header "BUILD SUMMARY"
    print_success "Build completed successfully"
    echo ""
    echo "Image tags:"
    echo "  - ${TAG_LATEST}"
    echo "  - ${TAG_COMMIT}"
    echo "  - ${TAG_TIMESTAMP}"
    echo ""

    if [ "$DO_PUSH" = true ] && docker info 2>/dev/null | grep -q "Username: ${DOCKER_REGISTRY}"; then
        echo "Images pushed to Docker Hub: https://hub.docker.com/r/${DOCKER_REGISTRY}/${IMAGE_NAME}"
    else
        echo "Images built locally (not pushed)"
    fi

    if [ -n "$BUILD_START_TIME" ] && [ -n "$BUILD_END_TIME" ]; then
        local total_time
        total_time=$(time_diff "$BUILD_START_TIME" "$BUILD_END_TIME")
        echo "Total build time: ${total_time}"
    fi
    echo ""

    print_success "Done!"
    exit 0
}

# Run main function
main "$@"