#!/bin/bash
# BuildKit 离线构建脚本 v2 - 增强版
# Usage: ./build-offline-v2.sh [cache-dir] [version]

set -e

CACHE_DIR="${1:-./build-cache}"
VERSION="${2:-latest}"
BUILDKIT_CACHE_DIR="$CACHE_DIR/buildkit"
IMAGES_DIR="$CACHE_DIR/images"
RESOURCES_DIR="$CACHE_DIR/resources"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

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

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

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

# 检查缓存目录
if [ ! -d "$CACHE_DIR" ]; then
    log_error "缓存目录 $CACHE_DIR 不存在"
    log_info "请先解压缓存包: tar -xzf build-cache-*.tar.gz"
    exit 1
fi

# 确保 buildx 构建器存在（使用 docker 驱动以支持本地镜像）
setup_buildx() {
    log_info "设置 BuildKit 构建器..."
    
    # 删除旧的构建器
    if docker buildx inspect offline-builder > /dev/null 2>&1; then
        docker buildx rm offline-builder 2>/dev/null || true
    fi
    
    # 创建新的构建器，使用 docker 驱动（支持本地镜像，不需要推送到 registry）
    docker buildx create --name offline-builder \
        --driver docker-container \
        --driver-opt image=moby/buildkit:buildx-stable-1 \
        --use
    
    log_info "BuildKit 构建器创建完成"
}

# 加载基础镜像
load_base_images() {
    log_info "加载基础镜像..."
    
    if [ ! -f "$IMAGES_DIR/base-images.tar" ]; then
        log_warn "基础镜像文件不存在: $IMAGES_DIR/base-images.tar"
        log_info "检查本地是否存在所需镜像..."
        
        # 检查关键镜像是否存在
        required_images=(
            "maven:3-eclipse-temurin-21"
            "eclipse-temurin:21-jdk"
            "mysql:8"
            "node:20-alpine"
            "nginx:1.29"
        )
        
        for img in "${required_images[@]}"; do
            if ! docker inspect "$img" > /dev/null 2>&1; then
                log_error "缺少基础镜像: $img"
                log_info "请确保基础镜像已加载: docker load -i base-images.tar"
                exit 1
            fi
        done
        
        log_info "本地基础镜像检查通过"
        return
    fi
    
    log_info "从 $IMAGES_DIR/base-images.tar 加载基础镜像..."
    docker load -i "$IMAGES_DIR/base-images.tar"
    log_info "基础镜像加载完成"
}

# 验证镜像是否存在
verify_image() {
    local image_name=$1
    if docker inspect "$image_name" > /dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# 离线构建函数
offline_build() {
    local service_name=$1
    local image_name=$2
    local dockerfile=$3
    local cache_file="$BUILDKIT_CACHE_DIR/${service_name}-cache"
    
    log_info "----------------------------------------"
    log_info "构建 [$service_name] -> $image_name:$VERSION"
    log_info "----------------------------------------"
    
    if [ ! -d "$cache_file" ]; then
        log_warn "$service_name 的缓存不存在，跳过..."
        return 0
    fi
    
    # 获取 Dockerfile 中的基础镜像
    local base_images
    base_images=$(grep -E '^FROM' "$dockerfile" | awk '{print $2}' | sort -u)
    
    log_info "检查基础镜像..."
    for base_img in $base_images; do
        # 跳过多阶段构建中的 AS 别名
        base_img=$(echo "$base_img" | cut -d':' -f1-2 | sed 's/AS.*//i' | tr -d ' ')
        
        if [ -z "$base_img" ] || [[ "$base_img" == *"AS"* ]]; then
            continue
        fi
        
        if verify_image "$base_img"; then
            log_info "  ✓ $base_img"
        else
            log_warn "  ✗ $base_img (未找到)"
            # 尝试从 base-images.tar 中加载
            if [ -f "$IMAGES_DIR/base-images.tar" ]; then
                log_info "  尝试从 tar 包加载..."
                docker load -i "$IMAGES_DIR/base-images.tar" 2>/dev/null || true
            fi
        fi
    done
    
    # 执行离线构建
    log_info "开始构建..."
    
    # 构建参数
    local build_args=()
    
    # 为需要外部资源的服务添加 build-arg
    case "$service_name" in
        runtime|deer-flow-backend|deer-flow-frontend)
            if [ -d "$RESOURCES_DIR" ]; then
                build_args+=("--build-arg" "RESOURCES_DIR=$RESOURCES_DIR")
            fi
            ;;
        backend-python)
            if [ -d "$RESOURCES_DIR" ]; then
                build_args+=("--build-arg" "RESOURCES_DIR=$RESOURCES_DIR")
                build_args+=("--build-arg" "DATAX_LOCAL_PATH=$RESOURCES_DIR/DataX")
            fi
            ;;
    esac
    
    # 执行构建
    if docker buildx build \
        --builder offline-builder \
        --cache-from "type=local,src=$cache_file" \
        --pull=false \
        --output "type=docker" \
        "${build_args[@]}" \
        -f "$dockerfile" \
        -t "$image_name:$VERSION" \
        . 2>&1; then
        
        log_info "✓ $service_name 构建成功"
        return 0
    else
        log_error "✗ $service_name 构建失败"
        return 1
    fi
}

# 主流程
main() {
    log_info "======================================"
    log_info "BuildKit 离线构建"
    log_info "======================================"
    log_info "缓存目录: $CACHE_DIR"
    log_info "版本: $VERSION"
    
    # 步骤 1: 设置构建器
    setup_buildx
    
    # 步骤 2: 加载基础镜像
    load_base_images
    
    # 步骤 3: 定义服务列表
    declare -A SERVICES=(
        ["database"]="datamate-database:scripts/images/database/Dockerfile"
        ["gateway"]="datamate-gateway:scripts/images/gateway/Dockerfile"
        ["backend"]="datamate-backend:scripts/images/backend/Dockerfile"
        ["frontend"]="datamate-frontend:scripts/images/frontend/Dockerfile"
        ["runtime"]="datamate-runtime:scripts/images/runtime/Dockerfile"
        ["backend-python"]="datamate-backend-python:scripts/images/backend-python/Dockerfile"
        ["deer-flow-backend"]="deer-flow-backend:scripts/images/deer-flow-backend/Dockerfile"
        ["deer-flow-frontend"]="deer-flow-frontend:scripts/images/deer-flow-frontend/Dockerfile"
        ["mineru"]="datamate-mineru:scripts/images/mineru/Dockerfile"
    )
    
    # 步骤 4: 批量构建
    log_info ""
    log_info "======================================"
    log_info "开始批量构建"
    log_info "======================================"
    
    local failed=()
    local succeeded=()
    
    for service_name in "${!SERVICES[@]}"; do
        IFS=':' read -r image_name dockerfile <<< "${SERVICES[$service_name]}"
        
        if offline_build "$service_name" "$image_name" "$dockerfile"; then
            succeeded+=("$service_name")
        else
            failed+=("$service_name")
        fi
        echo ""
    done
    
    # 步骤 5: 汇总结果
    log_info "======================================"
    log_info "构建完成"
    log_info "======================================"
    
    if [ ${#succeeded[@]} -gt 0 ]; then
        log_info "成功 (${#succeeded[@]}): ${succeeded[*]}"
    fi
    
    if [ ${#failed[@]} -gt 0 ]; then
        log_error "失败 (${#failed[@]}): ${failed[*]}"
        exit 1
    else
        log_info "✓ 所有服务构建成功！"
        echo ""
        log_info "镜像列表:"
        docker images --format "table {{.Repository}}:{{.Tag}}\t{{.Size}}" | grep -E "(datamate-|deer-flow-)" || true
    fi
}

# 执行主流程
main "$@"