#!/bin/bash

# Docker 镜像管理脚本
# 支持构建、保存、加载、清理镜像

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TEST_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
cd "$TEST_DIR"

# 检测 docker-compose 命令
if command -v docker-compose &> /dev/null; then
    DOCKER_COMPOSE="docker-compose"
elif docker compose version &> /dev/null 2>&1; then
    DOCKER_COMPOSE="docker compose"
else
    echo "错误: 未找到 docker-compose 或 docker compose 命令"
    exit 1
fi

# 镜像缓存目录
IMAGE_CACHE_DIR="$TEST_DIR/images-cache"
mkdir -p "$IMAGE_CACHE_DIR"

# 定义镜像列表
IMAGES=(
    "argus-metric-ftp:latest"
    "argus-metric-prometheus:latest"
    "argus-metric-grafana:latest"
)

# 镜像文件名映射
declare -A IMAGE_FILES=(
    ["argus-metric-ftp:latest"]="argus-ftp.tar"
    ["argus-metric-prometheus:latest"]="argus-prometheus.tar"
    ["argus-metric-grafana:latest"]="argus-grafana.tar"
)

# 检查镜像是否存在
check_image_exists() {
    local image=$1
    if docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "^${image}$"; then
        return 0
    else
        return 1
    fi
}

# 加载镜像
load_image() {
    local image=$1
    local file="${IMAGE_CACHE_DIR}/${IMAGE_FILES[$image]}"
    
    if [ -f "$file" ]; then
        echo "正在从缓存加载镜像: $image"
        docker load -i "$file"
        return 0
    else
        return 1
    fi
}

# 保存镜像
save_image() {
    local image=$1
    local file="${IMAGE_CACHE_DIR}/${IMAGE_FILES[$image]}"
    
    if check_image_exists "$image"; then
        echo "正在保存镜像到缓存: $image"
        docker save -o "$file" "$image"
        echo "已保存: $file ($(du -h "$file" | cut -f1))"
        return 0
    else
        echo "镜像不存在: $image"
        return 1
    fi
}

# 构建所有镜像
build_all() {
    echo "=========================================="
    echo "  构建所有 Docker 镜像"
    echo "=========================================="
    echo ""
    
    local build_flag="${1:---no-cache}"
    
    echo "开始构建镜像..."
    $DOCKER_COMPOSE build $build_flag
    
    echo ""
    echo "构建完成！"
}

# 保存所有镜像
save_all() {
    echo "=========================================="
    echo "  保存所有 Docker 镜像到缓存"
    echo "=========================================="
    echo ""
    
    for image in "${IMAGES[@]}"; do
        if save_image "$image"; then
            echo "✓ $image"
        else
            echo "✗ $image (跳过)"
        fi
        echo ""
    done
    
    echo "缓存目录: $IMAGE_CACHE_DIR"
    echo "总大小: $(du -sh "$IMAGE_CACHE_DIR" | cut -f1)"
}

# 加载所有镜像
load_all() {
    echo "=========================================="
    echo "  从缓存加载所有 Docker 镜像"
    echo "=========================================="
    echo ""
    
    local loaded=0
    local skipped=0
    
    for image in "${IMAGES[@]}"; do
        if check_image_exists "$image"; then
            echo "镜像已存在，跳过: $image"
            ((skipped++))
        elif load_image "$image"; then
            echo "✓ 已加载: $image"
            ((loaded++))
        else
            echo "✗ 缓存不存在: $image"
        fi
        echo ""
    done
    
    echo "加载: $loaded, 跳过: $skipped"
}

# 检查镜像状态
status() {
    echo "=========================================="
    echo "  镜像状态"
    echo "=========================================="
    echo ""
    
    echo "Docker 镜像:"
    for image in "${IMAGES[@]}"; do
        if check_image_exists "$image"; then
            local size=$(docker images --format "{{.Size}}" "$image" | head -1)
            echo "  ✓ $image ($size)"
        else
            echo "  ✗ $image (未构建)"
        fi
    done
    
    echo ""
    echo "缓存文件:"
    if [ -d "$IMAGE_CACHE_DIR" ] && [ "$(ls -A $IMAGE_CACHE_DIR 2>/dev/null)" ]; then
        for image in "${IMAGES[@]}"; do
            local file="${IMAGE_CACHE_DIR}/${IMAGE_FILES[$image]}"
            if [ -f "$file" ]; then
                echo "  ✓ ${IMAGE_FILES[$image]} ($(du -h "$file" | cut -f1))"
            else
                echo "  ✗ ${IMAGE_FILES[$image]} (不存在)"
            fi
        done
        echo ""
        echo "缓存总大小: $(du -sh "$IMAGE_CACHE_DIR" | cut -f1)"
    else
        echo "  (无缓存文件)"
    fi
}

# 清理缓存
clean_cache() {
    echo "=========================================="
    echo "  清理镜像缓存"
    echo "=========================================="
    echo ""
    
    if [ -d "$IMAGE_CACHE_DIR" ] && [ "$(ls -A $IMAGE_CACHE_DIR 2>/dev/null)" ]; then
        echo "缓存目录: $IMAGE_CACHE_DIR"
        echo "大小: $(du -sh "$IMAGE_CACHE_DIR" | cut -f1)"
        echo ""
        read -p "确认删除所有缓存文件? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -rf "$IMAGE_CACHE_DIR"/*.tar
            echo "已清理缓存文件"
        else
            echo "已取消"
        fi
    else
        echo "没有缓存文件"
    fi
}

# 清理 Docker 镜像
clean_images() {
    echo "=========================================="
    echo "  清理 Docker 镜像"
    echo "=========================================="
    echo ""
    
    local exists=0
    for image in "${IMAGES[@]}"; do
        if check_image_exists "$image"; then
            exists=1
            break
        fi
    done
    
    if [ $exists -eq 0 ]; then
        echo "没有需要清理的镜像"
        return
    fi
    
    echo "将删除以下镜像:"
    for image in "${IMAGES[@]}"; do
        if check_image_exists "$image"; then
            echo "  - $image"
        fi
    done
    echo ""
    
    read -p "确认删除这些镜像? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        for image in "${IMAGES[@]}"; do
            if check_image_exists "$image"; then
                docker rmi "$image"
                echo "已删除: $image"
            fi
        done
    else
        echo "已取消"
    fi
}

# 智能准备镜像（自动检测并加载或构建）
prepare() {
    echo "=========================================="
    echo "  智能准备 Docker 镜像"
    echo "=========================================="
    echo ""
    
    local need_build=()
    local loaded=0
    local existed=0
    
    for image in "${IMAGES[@]}"; do
        if check_image_exists "$image"; then
            echo "✓ 镜像已存在: $image"
            ((existed++))
        elif load_image "$image"; then
            echo "✓ 已从缓存加载: $image"
            ((loaded++))
        else
            echo "✗ 需要构建: $image"
            need_build+=("$image")
        fi
    done
    
    echo ""
    echo "统计: 已存在 $existed, 已加载 $loaded, 需构建 ${#need_build[@]}"
    
    if [ ${#need_build[@]} -gt 0 ]; then
        echo ""
        echo "需要构建以下镜像:"
        for image in "${need_build[@]}"; do
            echo "  - $image"
        done
        echo ""
        
        read -p "是否现在构建? (Y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Nn]$ ]]; then
            build_all ""
            echo ""
            read -p "是否保存新构建的镜像到缓存? (Y/n): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Nn]$ ]]; then
                save_all
            fi
        fi
    else
        echo ""
        echo "所有镜像已就绪！"
    fi
}

# 显示帮助
show_help() {
    cat << EOF
Docker 镜像管理工具

用法: $0 <command>

命令:
  prepare          智能准备镜像（推荐）- 自动检测、加载或构建
  build            构建所有镜像
  build-cache      使用缓存构建
  save             保存所有镜像到缓存
  load             从缓存加载所有镜像
  status           查看镜像状态
  clean-cache      清理缓存文件
  clean-images     清理 Docker 镜像
  clean-all        清理缓存和镜像
  help             显示此帮助信息

示例:
  # 智能准备（首次使用或镜像丢失时）
  $0 prepare
  
  # 构建并保存镜像
  $0 build
  $0 save
  
  # 从缓存加载镜像
  $0 load
  
  # 查看状态
  $0 status

镜像缓存目录: $IMAGE_CACHE_DIR/
EOF
}

# 主逻辑
case "${1:-help}" in
    prepare)
        prepare
        ;;
    build)
        build_all "--no-cache"
        ;;
    build-cache)
        build_all ""
        ;;
    save)
        save_all
        ;;
    load)
        load_all
        ;;
    status)
        status
        ;;
    clean-cache)
        clean_cache
        ;;
    clean-images)
        clean_images
        ;;
    clean-all)
        clean_cache
        clean_images
        ;;
    help|--help|-h)
        show_help
        ;;
    *)
        echo "错误: 未知命令 '$1'"
        echo ""
        show_help
        exit 1
        ;;
esac