#!/bin/bash

# Docker镜像导入和验证脚本
# 用于导入Docker镜像并验证其完整性和兼容性

set -e

# 脚本配置
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEFAULT_IMAGES_DIR="$(dirname "$SCRIPT_DIR")/docker-images"
LOG_FILE="${DEFAULT_IMAGES_DIR}/import-verify.log"

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

# 日志函数
log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1" | tee -a "$LOG_FILE"
}

log_success() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] ✓${NC} $1" | tee -a "$LOG_FILE"
}

log_warning() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] ⚠${NC} $1" | tee -a "$LOG_FILE"
}

log_error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ✗${NC} $1" | tee -a "$LOG_FILE"
}

# 显示帮助信息
show_help() {
    cat << EOF
Docker镜像导入和验证脚本

用法: $0 [选项] [镜像目录]

选项:
    -h, --help              显示此帮助信息
    -v, --verify-only       仅验证，不导入
    -f, --force             强制导入，覆盖现有镜像
    -c, --check-manifest    检查清单文件
    --skip-compatibility    跳过兼容性检查
    --parallel N            并行导入数量 (默认: 2)
    --test-run              导入后运行测试容器

参数:
    镜像目录                包含Docker镜像文件的目录 (默认: ./docker-images)

示例:
    $0                                      # 导入默认目录中的所有镜像
    $0 -v /path/to/images                   # 仅验证镜像文件
    $0 -f --test-run                        # 强制导入并测试
    $0 --check-manifest                     # 检查清单文件完整性

EOF
}

# 默认配置
VERIFY_ONLY=false
FORCE_IMPORT=false
CHECK_MANIFEST=false
SKIP_COMPATIBILITY=false
PARALLEL_JOBS=2
TEST_RUN=false
IMAGES_DIR="$DEFAULT_IMAGES_DIR"

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            show_help
            exit 0
            ;;
        -v|--verify-only)
            VERIFY_ONLY=true
            shift
            ;;
        -f|--force)
            FORCE_IMPORT=true
            shift
            ;;
        -c|--check-manifest)
            CHECK_MANIFEST=true
            shift
            ;;
        --skip-compatibility)
            SKIP_COMPATIBILITY=true
            shift
            ;;
        --parallel)
            PARALLEL_JOBS="$2"
            if [[ ! "$PARALLEL_JOBS" =~ ^[1-9][0-9]*$ ]]; then
                log_error "并行任务数必须是正整数"
                exit 1
            fi
            shift 2
            ;;
        --test-run)
            TEST_RUN=true
            shift
            ;;
        -*)
            log_error "未知参数: $1"
            show_help
            exit 1
            ;;
        *)
            IMAGES_DIR="$1"
            shift
            ;;
    esac
done

# 检查输入目录
if [[ ! -d "$IMAGES_DIR" ]]; then
    log_error "镜像目录不存在: $IMAGES_DIR"
    exit 1
fi

# 创建日志目录
mkdir -p "$(dirname "$LOG_FILE")"

log "开始Docker镜像导入和验证..."
log "镜像目录: $IMAGES_DIR"
log "仅验证: $VERIFY_ONLY"
log "强制导入: $FORCE_IMPORT"
log "并行任务: $PARALLEL_JOBS"

# 检查Docker是否运行
check_docker() {
    log "检查Docker环境..."
    
    if ! command -v docker >/dev/null 2>&1; then
        log_error "Docker未安装"
        exit 1
    fi
    
    if ! docker info >/dev/null 2>&1; then
        log_error "Docker未运行或无法访问"
        exit 1
    fi
    
    local docker_version=$(docker --version)
    log_success "Docker环境正常: $docker_version"
    
    # 检查系统架构
    local system_arch=$(uname -m)
    log "系统架构: $system_arch"
    
    if [[ "$system_arch" != "x86_64" ]] && [[ "$SKIP_COMPATIBILITY" == false ]]; then
        log_warning "系统架构不是x86_64，可能存在兼容性问题"
    fi
}

# 检查清单文件
check_manifest_file() {
    if [[ "$CHECK_MANIFEST" == false ]]; then
        return 0
    fi
    
    local manifest_file="${IMAGES_DIR}/images-manifest.txt"
    
    log "检查清单文件..."
    
    if [[ ! -f "$manifest_file" ]]; then
        log_warning "清单文件不存在: $manifest_file"
        return 1
    fi
    
    log "验证清单文件中的镜像..."
    
    # 解析清单文件
    local current_file=""
    local expected_checksum=""
    local verification_failed=0
    
    while IFS= read -r line; do
        if [[ "$line" =~ ^文件:\ (.+)$ ]]; then
            current_file="${BASH_REMATCH[1]}"
        elif [[ "$line" =~ ^SHA256:\ (.+)$ ]]; then
            expected_checksum="${BASH_REMATCH[1]}"
            
            if [[ -n "$current_file" ]] && [[ -n "$expected_checksum" ]]; then
                local file_path="${IMAGES_DIR}/${current_file}"
                
                if [[ -f "$file_path" ]]; then
                    log "验证文件: $current_file"
                    local actual_checksum=$(sha256sum "$file_path" | cut -d' ' -f1)
                    
                    if [[ "$actual_checksum" == "$expected_checksum" ]]; then
                        log_success "校验和匹配: $current_file"
                    else
                        log_error "校验和不匹配: $current_file"
                        log_error "期望: $expected_checksum"
                        log_error "实际: $actual_checksum"
                        ((verification_failed++))
                    fi
                else
                    log_error "文件不存在: $current_file"
                    ((verification_failed++))
                fi
                
                current_file=""
                expected_checksum=""
            fi
        fi
    done < "$manifest_file"
    
    if [[ $verification_failed -eq 0 ]]; then
        log_success "清单文件验证通过"
        return 0
    else
        log_error "清单文件验证失败，$verification_failed 个文件有问题"
        return 1
    fi
}

# 验证镜像文件
verify_image_file() {
    local file="$1"
    local filename=$(basename "$file")
    
    log "验证镜像文件: $filename"
    
    # 检查文件是否存在
    if [[ ! -f "$file" ]]; then
        log_error "文件不存在: $filename"
        return 1
    fi
    
    # 检查文件大小
    local file_size=$(du -h "$file" | cut -f1)
    log "文件大小: $file_size"
    
    # 验证文件完整性
    if [[ "$file" == *.tar.gz ]]; then
        # 验证gzip文件
        if gzip -t "$file" 2>/dev/null; then
            log_success "压缩文件完整性验证通过: $filename"
        else
            log_error "压缩文件完整性验证失败: $filename"
            return 1
        fi
    elif [[ "$file" == *.tar ]]; then
        # 验证tar文件
        if tar -tf "$file" >/dev/null 2>&1; then
            log_success "tar文件完整性验证通过: $filename"
        else
            log_error "tar文件完整性验证失败: $filename"
            return 1
        fi
    else
        log_warning "未知文件类型，跳过验证: $filename"
        return 1
    fi
    
    return 0
}

# 导入镜像文件
import_image_file() {
    local file="$1"
    local filename=$(basename "$file")
    
    log "导入镜像文件: $filename"
    
    # 首先验证文件
    if ! verify_image_file "$file"; then
        log_error "文件验证失败，跳过导入: $filename"
        return 1
    fi
    
    # 检查是否需要强制导入
    local import_args=""
    if [[ "$FORCE_IMPORT" == true ]]; then
        import_args="--quiet"
    fi
    
    # 导入镜像
    local import_output
    if [[ "$file" == *.tar.gz ]]; then
        # 解压并导入
        import_output=$(gunzip -c "$file" | docker load 2>&1)
    else
        # 直接导入
        import_output=$(docker load -i "$file" 2>&1)
    fi
    
    if [[ $? -eq 0 ]]; then
        log_success "镜像导入成功: $filename"
        
        # 提取导入的镜像名称
        local imported_image=$(echo "$import_output" | grep "Loaded image" | sed 's/Loaded image: //')
        if [[ -n "$imported_image" ]]; then
            log "导入的镜像: $imported_image"
            
            # 验证镜像架构
            if [[ "$SKIP_COMPATIBILITY" == false ]]; then
                verify_image_architecture "$imported_image"
            fi
        fi
        
        return 0
    else
        log_error "镜像导入失败: $filename"
        log_error "错误信息: $import_output"
        return 1
    fi
}

# 验证镜像架构
verify_image_architecture() {
    local image="$1"
    
    log "验证镜像架构: $image"
    
    # 获取镜像信息
    local image_info=$(docker image inspect "$image" 2>/dev/null)
    
    if [[ $? -ne 0 ]]; then
        log_error "无法获取镜像信息: $image"
        return 1
    fi
    
    # 提取架构信息
    local architecture=$(echo "$image_info" | grep -o '"Architecture":"[^"]*"' | cut -d'"' -f4)
    local os=$(echo "$image_info" | grep -o '"Os":"[^"]*"' | cut -d'"' -f4)
    
    log "镜像架构: $os/$architecture"
    
    # 检查架构兼容性
    local system_arch=$(uname -m)
    local expected_arch="amd64"
    
    if [[ "$system_arch" == "x86_64" ]]; then
        expected_arch="amd64"
    elif [[ "$system_arch" == "aarch64" ]]; then
        expected_arch="arm64"
    fi
    
    if [[ "$architecture" == "$expected_arch" ]]; then
        log_success "镜像架构兼容: $architecture"
        return 0
    else
        log_warning "镜像架构可能不兼容: $architecture (系统: $system_arch)"
        return 1
    fi
}

# 测试镜像运行
test_image_run() {
    local image="$1"
    
    log "测试镜像运行: $image"
    
    # 根据镜像类型选择测试命令
    local test_command=""
    local container_name="test-$(echo "$image" | sed 's/[\/:]/_/g')-$$"
    
    case "$image" in
        *mysql*)
            test_command="docker run --rm --name $container_name -e MYSQL_ROOT_PASSWORD=test -d $image"
            ;;
        *redis*)
            test_command="docker run --rm --name $container_name -d $image"
            ;;
        *meilisearch*)
            test_command="docker run --rm --name $container_name -e MEILI_MASTER_KEY=test -d $image"
            ;;
        *knowledge-base-app*)
            # 应用镜像需要更复杂的测试
            log_warning "应用镜像测试需要完整环境，跳过单独测试"
            return 0
            ;;
        *)
            log_warning "未知镜像类型，跳过运行测试: $image"
            return 0
            ;;
    esac
    
    # 运行测试容器
    if eval "$test_command"; then
        log "测试容器启动成功: $container_name"
        
        # 等待容器启动
        sleep 5
        
        # 检查容器状态
        if docker ps | grep -q "$container_name"; then
            log_success "镜像运行测试通过: $image"
            
            # 停止测试容器
            docker stop "$container_name" >/dev/null 2>&1 || true
            return 0
        else
            log_error "测试容器启动失败: $image"
            
            # 显示容器日志
            docker logs "$container_name" 2>/dev/null || true
            docker rm "$container_name" >/dev/null 2>&1 || true
            return 1
        fi
    else
        log_error "无法启动测试容器: $image"
        return 1
    fi
}

# 处理单个镜像文件
process_image_file() {
    local file="$1"
    local filename=$(basename "$file")
    
    log "处理镜像文件: $filename"
    
    if [[ "$VERIFY_ONLY" == true ]]; then
        verify_image_file "$file"
        return $?
    else
        if import_image_file "$file"; then
            # 如果需要测试运行
            if [[ "$TEST_RUN" == true ]]; then
                # 提取镜像名称进行测试
                local image_name=$(echo "$filename" | sed 's/\.tar.*$//' | sed 's/_/:/g')
                test_image_run "$image_name" || true
            fi
            return 0
        else
            return 1
        fi
    fi
}

# 主处理流程
main() {
    check_docker
    check_manifest_file
    
    # 查找镜像文件
    local image_files=($(find "$IMAGES_DIR" -name "*.tar*" -type f))
    
    if [[ ${#image_files[@]} -eq 0 ]]; then
        log_error "在 $IMAGES_DIR 中没有找到镜像文件"
        exit 1
    fi
    
    log "找到 ${#image_files[@]} 个镜像文件"
    
    # 处理镜像文件
    local processed=0
    local failed=0
    local total=${#image_files[@]}
    
    # 使用并行处理
    export -f process_image_file verify_image_file import_image_file verify_image_architecture test_image_run
    export -f log log_success log_warning log_error
    export VERIFY_ONLY FORCE_IMPORT SKIP_COMPATIBILITY TEST_RUN LOG_FILE
    export RED GREEN YELLOW BLUE NC
    
    for file in "${image_files[@]}"; do
        if process_image_file "$file"; then
            ((processed++))
        else
            ((failed++))
        fi
    done
    
    # 显示导入的镜像
    if [[ "$VERIFY_ONLY" == false ]]; then
        log "当前Docker镜像列表:"
        docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}\t{{.CreatedAt}}"
    fi
    
    # 显示总结
    log_success "处理完成！"
    log "总文件数: $total"
    log "成功处理: $processed"
    log "失败数量: $failed"
    
    if [[ $failed -gt 0 ]]; then
        log_warning "有 $failed 个文件处理失败，请检查日志"
        exit 1
    else
        log_success "所有文件处理成功"
    fi
}

# 执行主流程
main