#!/bin/bash

# Docker镜像压缩和完整性检查脚本
# 用于压缩导出的镜像文件并验证完整性

set -e

# 脚本配置
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
DEFAULT_INPUT_DIR="${PROJECT_ROOT}/docker-images"
LOG_FILE="${DEFAULT_INPUT_DIR}/compress-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              显示此帮助信息
    -c, --compress-level N  压缩级别 (1-9, 默认: 6)
    -k, --keep-original     保留原始文件
    -v, --verify-only       仅验证，不压缩
    -u, --uncompress        解压缩文件
    --parallel N            并行处理数量 (默认: 2)

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

示例:
    $0                                      # 压缩默认目录中的所有tar文件
    $0 -c 9 -k /path/to/images             # 最高压缩级别，保留原文件
    $0 --verify-only                       # 仅验证现有文件完整性
    $0 --uncompress                        # 解压缩所有.gz文件

EOF
}

# 默认配置
COMPRESS_LEVEL=6
KEEP_ORIGINAL=false
VERIFY_ONLY=false
UNCOMPRESS=false
PARALLEL_JOBS=2
INPUT_DIR="$DEFAULT_INPUT_DIR"

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            show_help
            exit 0
            ;;
        -c|--compress-level)
            COMPRESS_LEVEL="$2"
            if [[ ! "$COMPRESS_LEVEL" =~ ^[1-9]$ ]]; then
                log_error "压缩级别必须是1-9之间的数字"
                exit 1
            fi
            shift 2
            ;;
        -k|--keep-original)
            KEEP_ORIGINAL=true
            shift
            ;;
        -v|--verify-only)
            VERIFY_ONLY=true
            shift
            ;;
        -u|--uncompress)
            UNCOMPRESS=true
            shift
            ;;
        --parallel)
            PARALLEL_JOBS="$2"
            if [[ ! "$PARALLEL_JOBS" =~ ^[1-9][0-9]*$ ]]; then
                log_error "并行任务数必须是正整数"
                exit 1
            fi
            shift 2
            ;;
        -*)
            log_error "未知参数: $1"
            show_help
            exit 1
            ;;
        *)
            INPUT_DIR="$1"
            shift
            ;;
    esac
done

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

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

log "开始镜像压缩和完整性检查..."
log "输入目录: $INPUT_DIR"
log "压缩级别: $COMPRESS_LEVEL"
log "保留原文件: $KEEP_ORIGINAL"
log "仅验证: $VERIFY_ONLY"
log "解压缩: $UNCOMPRESS"
log "并行任务: $PARALLEL_JOBS"

# 检查必要工具
check_tools() {
    local tools=("gzip" "sha256sum" "tar")
    
    for tool in "${tools[@]}"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            log_error "缺少必要工具: $tool"
            exit 1
        fi
    done
}

check_tools

# 验证文件完整性
verify_file() {
    local file="$1"
    local filename=$(basename "$file")
    
    log "验证文件: $filename"
    
    if [[ "$file" == *.tar.gz ]]; then
        # 验证gzip文件
        if gzip -t "$file" 2>/dev/null; then
            log_success "压缩文件完整性验证通过: $filename"
            return 0
        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"
            return 0
        else
            log_error "tar文件完整性验证失败: $filename"
            return 1
        fi
    else
        log_warning "未知文件类型，跳过验证: $filename"
        return 1
    fi
}

# 压缩文件
compress_file() {
    local file="$1"
    local filename=$(basename "$file")
    local compressed_file="${file}.gz"
    
    log "压缩文件: $filename (级别: $COMPRESS_LEVEL)"
    
    # 检查是否已经压缩
    if [[ "$file" == *.gz ]]; then
        log_warning "文件已经压缩，跳过: $filename"
        return 0
    fi
    
    # 检查压缩文件是否已存在
    if [[ -f "$compressed_file" ]]; then
        log_warning "压缩文件已存在，跳过: ${filename}.gz"
        return 0
    fi
    
    # 获取原始文件大小
    local original_size=$(du -b "$file" | cut -f1)
    local original_size_human=$(du -h "$file" | cut -f1)
    
    # 压缩文件
    if gzip -"$COMPRESS_LEVEL" -c "$file" > "$compressed_file"; then
        # 获取压缩后文件大小
        local compressed_size=$(du -b "$compressed_file" | cut -f1)
        local compressed_size_human=$(du -h "$compressed_file" | cut -f1)
        
        # 计算压缩比
        local ratio=$(echo "scale=2; $compressed_size * 100 / $original_size" | bc -l 2>/dev/null || echo "N/A")
        
        log_success "文件压缩成功: ${filename}.gz"
        log "原始大小: $original_size_human"
        log "压缩后大小: $compressed_size_human"
        if [[ "$ratio" != "N/A" ]]; then
            log "压缩比: ${ratio}%"
        fi
        
        # 验证压缩文件
        if verify_file "$compressed_file"; then
            # 删除原文件（如果不保留）
            if [[ "$KEEP_ORIGINAL" == false ]]; then
                rm "$file"
                log "已删除原文件: $filename"
            fi
            return 0
        else
            log_error "压缩文件验证失败，删除压缩文件"
            rm -f "$compressed_file"
            return 1
        fi
    else
        log_error "文件压缩失败: $filename"
        return 1
    fi
}

# 解压缩文件
uncompress_file() {
    local file="$1"
    local filename=$(basename "$file")
    
    log "解压缩文件: $filename"
    
    # 检查是否是压缩文件
    if [[ "$file" != *.gz ]]; then
        log_warning "文件未压缩，跳过: $filename"
        return 0
    fi
    
    # 生成解压后的文件名
    local uncompressed_file="${file%.gz}"
    local uncompressed_filename=$(basename "$uncompressed_file")
    
    # 检查解压文件是否已存在
    if [[ -f "$uncompressed_file" ]]; then
        log_warning "解压文件已存在，跳过: $uncompressed_filename"
        return 0
    fi
    
    # 解压文件
    if gunzip -c "$file" > "$uncompressed_file"; then
        log_success "文件解压成功: $uncompressed_filename"
        
        # 验证解压文件
        if verify_file "$uncompressed_file"; then
            # 删除压缩文件（如果不保留）
            if [[ "$KEEP_ORIGINAL" == false ]]; then
                rm "$file"
                log "已删除压缩文件: $filename"
            fi
            return 0
        else
            log_error "解压文件验证失败，删除解压文件"
            rm -f "$uncompressed_file"
            return 1
        fi
    else
        log_error "文件解压失败: $filename"
        return 1
    fi
}

# 处理单个文件
process_file() {
    local file="$1"
    
    if [[ "$VERIFY_ONLY" == true ]]; then
        verify_file "$file"
    elif [[ "$UNCOMPRESS" == true ]]; then
        uncompress_file "$file"
    else
        compress_file "$file"
    fi
}

# 查找需要处理的文件
if [[ "$UNCOMPRESS" == true ]]; then
    FILES=($(find "$INPUT_DIR" -name "*.tar.gz" -type f))
    log "找到 ${#FILES[@]} 个压缩文件"
elif [[ "$VERIFY_ONLY" == true ]]; then
    FILES=($(find "$INPUT_DIR" -name "*.tar*" -type f))
    log "找到 ${#FILES[@]} 个文件需要验证"
else
    FILES=($(find "$INPUT_DIR" -name "*.tar" -type f))
    log "找到 ${#FILES[@]} 个tar文件需要压缩"
fi

if [[ ${#FILES[@]} -eq 0 ]]; then
    log_warning "没有找到需要处理的文件"
    exit 0
fi

# 处理文件
PROCESSED=0
FAILED=0
TOTAL=${#FILES[@]}

# 使用并行处理
export -f process_file verify_file compress_file uncompress_file log log_success log_warning log_error
export COMPRESS_LEVEL KEEP_ORIGINAL VERIFY_ONLY UNCOMPRESS LOG_FILE
export RED GREEN YELLOW BLUE NC

printf '%s\n' "${FILES[@]}" | xargs -n 1 -P "$PARALLEL_JOBS" -I {} bash -c 'process_file "$@"' _ {}

# 统计结果
for file in "${FILES[@]}"; do
    if [[ "$VERIFY_ONLY" == true ]]; then
        if verify_file "$file" >/dev/null 2>&1; then
            ((PROCESSED++))
        else
            ((FAILED++))
        fi
    elif [[ "$UNCOMPRESS" == true ]]; then
        uncompressed_file="${file%.gz}"
        if [[ -f "$uncompressed_file" ]] || [[ "$file" != *.gz ]]; then
            ((PROCESSED++))
        else
            ((FAILED++))
        fi
    else
        compressed_file="${file}.gz"
        if [[ -f "$compressed_file" ]] || [[ "$file" == *.gz ]]; then
            ((PROCESSED++))
        else
            ((FAILED++))
        fi
    fi
done

# 更新清单文件
if [[ "$VERIFY_ONLY" == false ]]; then
    manifest_file="${INPUT_DIR}/images-manifest.txt"
    if [[ -f "$manifest_file" ]]; then
        log "更新镜像清单..."
        
        # 备份原清单
        cp "$manifest_file" "${manifest_file}.backup"
        
        # 重新生成清单
        cat > "$manifest_file" << EOF
# Docker镜像清单
# 更新时间: $(date)
# 处理目录: $INPUT_DIR

EOF
        
        for file in "$INPUT_DIR"/*.tar*; do
            if [[ -f "$file" ]]; then
                filename=$(basename "$file")
                size=$(du -h "$file" | cut -f1)
                checksum=$(sha256sum "$file" | cut -d' ' -f1)
                
                echo "文件: $filename" >> "$manifest_file"
                echo "大小: $size" >> "$manifest_file"
                echo "SHA256: $checksum" >> "$manifest_file"
                echo "" >> "$manifest_file"
            fi
        done
        
        log_success "镜像清单已更新"
    fi
fi

# 显示总结
log_success "处理完成！"
log "总文件数: $TOTAL"
log "成功处理: $PROCESSED"
log "失败数量: $FAILED"

if [[ "$FAILED" -gt 0 ]]; then
    log_warning "有 $FAILED 个文件处理失败，请检查日志"
    exit 1
else
    log_success "所有文件处理成功"
fi