#!/bin/bash

#============================================================================
# FOXHUNT HFT TRADING SYSTEM - DOCKER IMAGE BUILD SCRIPT
#============================================================================
# Builds optimized Docker images for all services using multi-stage builds
#
# Usage:
#   ./build_docker_images.sh [OPTIONS]
#
# Options:
#   --service SERVICE     Build specific service (trading, ml, backtesting, tli, all)
#   --tag-prefix PREFIX   Docker image tag prefix (default: foxhunt)
#   --version VERSION     Version tag (default: latest)
#   --enable-gpu         Enable CUDA/GPU support
#   --target-arch ARCH   Target architecture (x86_64, aarch64)
#   --build-mode MODE    Build mode (release, debug)
#   --push               Push images to registry after building
#   --registry URL       Docker registry URL
#   --parallel           Build images in parallel
#   --help               Show this help message
#============================================================================

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
TAG_PREFIX="${TAG_PREFIX:-foxhunt}"
VERSION="${VERSION:-latest}"
ENABLE_GPU="${ENABLE_GPU:-false}"
TARGET_ARCH="${TARGET_ARCH:-x86_64}"
BUILD_MODE="${BUILD_MODE:-release}"
PUSH_IMAGES="${PUSH_IMAGES:-false}"
REGISTRY_URL="${REGISTRY_URL:-}"
BUILD_PARALLEL="${BUILD_PARALLEL:-false}"
SERVICE="${SERVICE:-all}"

# Service configurations
declare -A SERVICES=(
    ["trading"]="trading_service:8080"
    ["ml"]="ml_training_service:8082"
    ["backtesting"]="backtesting_service:8083"
    ["tli"]="tli:8081"
)

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

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

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

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

show_help() {
    cat << EOF
Foxhunt HFT Trading System - Docker Image Builder

Usage: $0 [OPTIONS]

OPTIONS:
    --service SERVICE     Build specific service (trading, ml, backtesting, tli, all)
    --tag-prefix PREFIX   Docker image tag prefix (default: foxhunt)
    --version VERSION     Version tag (default: latest)
    --enable-gpu         Enable CUDA/GPU support
    --target-arch ARCH   Target architecture (x86_64, aarch64)
    --build-mode MODE    Build mode (release, debug)
    --push               Push images to registry after building
    --registry URL       Docker registry URL
    --parallel           Build images in parallel
    --help               Show this help message

SERVICES:
    trading              Trading Service
    ml                   ML Training Service
    backtesting          Backtesting Service
    tli                  Terminal Line Interface
    all                  Build all services

EXAMPLES:
    $0                                    # Build all services
    $0 --service trading --push           # Build and push trading service
    $0 --enable-gpu --parallel            # Build all with GPU support in parallel
    $0 --target-arch aarch64              # Build for ARM64 architecture

EOF
}

# Parse arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --service)
                SERVICE="$2"
                shift 2
                ;;
            --tag-prefix)
                TAG_PREFIX="$2"
                shift 2
                ;;
            --version)
                VERSION="$2"
                shift 2
                ;;
            --enable-gpu)
                ENABLE_GPU=true
                shift
                ;;
            --target-arch)
                TARGET_ARCH="$2"
                shift 2
                ;;
            --build-mode)
                BUILD_MODE="$2"
                shift 2
                ;;
            --push)
                PUSH_IMAGES=true
                shift
                ;;
            --registry)
                REGISTRY_URL="$2"
                shift 2
                ;;
            --parallel)
                BUILD_PARALLEL=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

# Check prerequisites
check_prerequisites() {
    log_info "Checking Docker build prerequisites..."
    
    # Check Docker
    if ! command -v docker &> /dev/null; then
        log_error "Docker is not installed"
        exit 1
    fi
    
    # Check Docker version for multi-stage builds
    local docker_version=$(docker --version | grep -oE '[0-9]+\.[0-9]+' | head -1)
    local major_version=$(echo "$docker_version" | cut -d. -f1)
    local minor_version=$(echo "$docker_version" | cut -d. -f2)
    
    if [[ $major_version -lt 17 ]] || [[ $major_version -eq 17 && $minor_version -lt 5 ]]; then
        log_error "Docker version 17.05 or higher is required for multi-stage builds"
        exit 1
    fi
    
    # Check if Docker daemon is running
    if ! docker info &> /dev/null; then
        log_error "Docker daemon is not running"
        exit 1
    fi
    
    # Check Dockerfile exists
    local dockerfile="$SCRIPT_DIR/Dockerfile.production"
    if [[ ! -f "$dockerfile" ]]; then
        log_error "Production Dockerfile not found: $dockerfile"
        exit 1
    fi
    
    # Check if buildx is available for cross-platform builds
    if [[ "$TARGET_ARCH" != "x86_64" ]]; then
        if ! docker buildx version &> /dev/null; then
            log_warning "Docker Buildx not available - cross-platform builds may not work"
        fi
    fi
    
    log_success "Prerequisites check completed"
}

# Create Docker builder for multi-platform builds
create_builder() {
    if [[ "$TARGET_ARCH" == "x86_64" ]]; then
        return 0
    fi
    
    log_info "Creating Docker builder for multi-platform builds..."
    
    local builder_name="foxhunt-builder"
    
    # Check if builder already exists
    if ! docker buildx inspect "$builder_name" &> /dev/null; then
        docker buildx create --name "$builder_name" --driver docker-container --bootstrap
        log_success "Created Docker builder: $builder_name"
    else
        log_info "Using existing Docker builder: $builder_name"
    fi
    
    docker buildx use "$builder_name"
}

# Build Docker image for a service
build_service_image() {
    local service_name="$1"
    local service_config="${SERVICES[$service_name]}"
    
    if [[ -z "$service_config" ]]; then
        log_error "Unknown service: $service_name"
        return 1
    fi
    
    IFS=':' read -r binary_name port <<< "$service_config"
    
    local image_name="$TAG_PREFIX/$service_name"
    local full_tag="$image_name:$VERSION"
    
    # Add registry prefix if specified
    if [[ -n "$REGISTRY_URL" ]]; then
        full_tag="$REGISTRY_URL/$full_tag"
    fi
    
    log_info "Building Docker image: $full_tag"
    log_info "Service: $service_name, Binary: $binary_name, Target: $service_name"
    
    local build_start=$(date +%s)
    
    # Prepare build arguments
    local build_args=(
        "--build-arg" "SERVICE_NAME=$binary_name"
        "--build-arg" "TARGET_ARCH=$TARGET_ARCH"
        "--build-arg" "ENABLE_GPU=$ENABLE_GPU"
        "--build-arg" "BUILD_MODE=$BUILD_MODE"
        "--target" "$service_name"
        "--tag" "$full_tag"
        "--file" "$SCRIPT_DIR/Dockerfile.production"
    )
    
    # Add platform specification for cross-platform builds
    if [[ "$TARGET_ARCH" != "x86_64" ]]; then
        case "$TARGET_ARCH" in
            "aarch64")
                build_args+=("--platform" "linux/arm64")
                ;;
            *)
                log_error "Unsupported target architecture: $TARGET_ARCH"
                return 1
                ;;
        esac
    fi
    
    # Add context
    build_args+=("$PROJECT_ROOT")
    
    # Execute build
    if docker build "${build_args[@]}"; then
        local build_end=$(date +%s)
        local build_duration=$((build_end - build_start))
        
        # Get image size
        local image_size=$(docker images "$full_tag" --format "table {{.Size}}" | tail -n 1)
        
        log_success "$service_name image built successfully (${build_duration}s, ${image_size})"
        
        # Tag with additional tags if needed
        tag_additional_versions "$full_tag" "$service_name"
        
        # Push image if requested
        if [[ "$PUSH_IMAGES" == "true" ]]; then
            push_image "$full_tag"
        fi
        
        return 0
    else
        log_error "Failed to build $service_name image"
        return 1
    fi
}

# Tag image with additional versions
tag_additional_versions() {
    local full_tag="$1"
    local service_name="$2"
    
    local base_name="${full_tag%:*}"
    
    # Add git commit tag if in git repository
    if git rev-parse HEAD &>/dev/null; then
        local git_commit=$(git rev-parse --short HEAD)
        local commit_tag="$base_name:$git_commit"
        docker tag "$full_tag" "$commit_tag"
        log_info "Tagged with commit: $commit_tag"
    fi
    
    # Add architecture-specific tag
    local arch_tag="$base_name:$VERSION-$TARGET_ARCH"
    docker tag "$full_tag" "$arch_tag"
    log_info "Tagged with architecture: $arch_tag"
    
    # Add GPU-specific tag if enabled
    if [[ "$ENABLE_GPU" == "true" ]]; then
        local gpu_tag="$base_name:$VERSION-gpu"
        docker tag "$full_tag" "$gpu_tag"
        log_info "Tagged with GPU support: $gpu_tag"
    fi
}

# Push image to registry
push_image() {
    local image_tag="$1"
    
    log_info "Pushing image: $image_tag"
    
    if docker push "$image_tag"; then
        log_success "Successfully pushed: $image_tag"
        
        # Push additional tags
        local base_name="${image_tag%:*}"
        
        # Push commit tag
        if git rev-parse HEAD &>/dev/null; then
            local git_commit=$(git rev-parse --short HEAD)
            docker push "$base_name:$git_commit" 2>/dev/null || true
        fi
        
        # Push architecture tag
        docker push "$base_name:$VERSION-$TARGET_ARCH" 2>/dev/null || true
        
        # Push GPU tag if enabled
        if [[ "$ENABLE_GPU" == "true" ]]; then
            docker push "$base_name:$VERSION-gpu" 2>/dev/null || true
        fi
    else
        log_error "Failed to push: $image_tag"
        return 1
    fi
}

# Build all services in parallel
build_all_parallel() {
    log_info "Building all services in parallel..."
    
    local pids=()
    local failed_services=()
    
    for service in "${!SERVICES[@]}"; do
        (
            if build_service_image "$service"; then
                echo "$service: SUCCESS"
            else
                echo "$service: FAILED"
            fi
        ) &
        pids+=($!)
    done
    
    # Wait for all builds to complete
    for i in "${!pids[@]}"; do
        local pid=${pids[$i]}
        local service=${!SERVICES[@]:$i:1}
        
        if wait $pid; then
            log_success "$service build completed"
        else
            log_error "$service build failed"
            failed_services+=("$service")
        fi
    done
    
    if [[ ${#failed_services[@]} -eq 0 ]]; then
        log_success "All services built successfully"
        return 0
    else
        log_error "Failed to build services: ${failed_services[*]}"
        return 1
    fi
}

# Build all services sequentially
build_all_sequential() {
    log_info "Building all services sequentially..."
    
    local failed_services=()
    
    for service in "${!SERVICES[@]}"; do
        if build_service_image "$service"; then
            log_success "$service build completed"
        else
            log_error "$service build failed"
            failed_services+=("$service")
        fi
    done
    
    if [[ ${#failed_services[@]} -eq 0 ]]; then
        log_success "All services built successfully"
        return 0
    else
        log_error "Failed to build services: ${failed_services[*]}"
        return 1
    fi
}

# Generate build summary
generate_summary() {
    log_info "Generating build summary..."
    
    local total_size=0
    
    cat << EOF

${GREEN}=============================================================================
FOXHUNT HFT TRADING SYSTEM - DOCKER BUILD SUMMARY
=============================================================================${NC}

${BLUE}Build Configuration:${NC}
  • Tag Prefix:         $TAG_PREFIX
  • Version:            $VERSION
  • Target Architecture: $TARGET_ARCH
  • GPU Support:        $ENABLE_GPU
  • Build Mode:         $BUILD_MODE
  • Registry:           ${REGISTRY_URL:-"(local)"}

${BLUE}Built Images:${NC}
EOF

    for service in "${!SERVICES[@]}"; do
        local image_name="$TAG_PREFIX/$service:$VERSION"
        if [[ -n "$REGISTRY_URL" ]]; then
            image_name="$REGISTRY_URL/$image_name"
        fi
        
        if docker images "$image_name" --format "table {{.Repository}}:{{.Tag}} {{.Size}}" | tail -n 1 | grep -q "$service"; then
            local size=$(docker images "$image_name" --format "table {{.Size}}" | tail -n 1)
            echo "  • $service: $image_name ($size)"
        fi
    done

    cat << EOF

${BLUE}Management Commands:${NC}
  • Run service:        docker run -d $TAG_PREFIX/<service>:$VERSION
  • View images:        docker images | grep $TAG_PREFIX
  • Remove images:      docker rmi $TAG_PREFIX/<service>:$VERSION
  • Deploy:             docker-compose -f docker-compose.production.yml up -d

${GREEN}Docker image build completed successfully!${NC}

EOF
}

# Main build logic
main() {
    parse_args "$@"
    
    log_info "Starting Foxhunt Docker image build..."
    log_info "Service: $SERVICE, Version: $VERSION, Architecture: $TARGET_ARCH"
    
    check_prerequisites
    create_builder
    
    case "$SERVICE" in
        "all")
            if [[ "$BUILD_PARALLEL" == "true" ]]; then
                build_all_parallel
            else
                build_all_sequential
            fi
            ;;
        "trading"|"ml"|"backtesting"|"tli")
            build_service_image "$SERVICE"
            ;;
        *)
            log_error "Unknown service: $SERVICE"
            log_info "Available services: ${!SERVICES[*]} all"
            exit 1
            ;;
    esac
    
    generate_summary
    
    log_success "Foxhunt Docker image build completed!"
}

# Handle interruption
trap 'log_error "Build interrupted"; exit 1' INT TERM

# Run main function
main "$@"