#!/bin/bash

TAG_PREFIX=

# 保存镜像的函数
save() {
    # 如果images目录不存在，则创建它
    if [ ! -d "./images" ]; then
        mkdir images
    fi
    cd images

    # 获取带有指定前缀的镜像列表并保存到文件
    docker images --format "{{.ID}} {{.Repository}}:{{.Tag}} {{.Size}}" | grep "$TAG_PREFIX" > images_list.txt

    # 遍历每条镜像记录
    while read -r line; do
        image_id=$(echo $line | awk '{print $1}')
        # 更安全地处理Repository和Tag，特别是当Tag包含特殊字符时
        image_parts=($(echo $line | awk '{split($0,a," "); print a[2]}' | tr ":" "\n"))
        image_repository=${image_parts[0]}
        image_tag=${image_parts[1]:-${TAG_PREFIX}-latest} # 如果没有Tag，默认给一个值，避免空Tag问题
        image_size=$(echo $line | awk '{print $3}')
    
        if [[ ! -f "$image_id.tar" ]]; then
            # 确保Tag格式正确，特别处理无Tag情况
            if [ -z "$image_tag" ]; then
                image_load_cmd="$image_repository"
            else
                image_load_cmd="$image_repository:$image_tag"
            fi
            docker save -o "$image_id.tar" "$image_load_cmd" && \
            echo "镜像 ${image_load_cmd} 已保存，大小 $image_size"
        else
            echo "ID为$image_id的镜像已保存，跳过..."
            echo "$image_id $image_repository:$image_tag" >> existing_tags.txt
        fi
    done < images_list.txt
}

# 加载镜像的函数
load() {
    cd images

    # 先加载所有镜像文件
    while read -r line; do
        image_id=$(echo $line | awk '{print $1}')
        if [[ -f "$image_id.tar" ]]; then
            docker load -i $image_id.tar && \
            echo "ID为$image_id的镜像已加载"
        fi
    done < images_list.txt

    # 根据记录恢复所有标签
    while read -r line; do
        image_id=$(echo $line | awk '{print $1}')
        image_repo_tag=$(echo $line | awk '{print $2}')
        if docker images | grep -q $image_id; then
            docker tag $image_id $image_repo_tag && \
            echo "为ID为$image_id的镜像恢复了标签 $image_repo_tag"
        else
            echo "未找到ID为$image_id的镜像进行标签恢复。"
        fi
    done < existing_tags.txt
}

# 脚本主逻辑
if [ -z "$1" ]; then
    echo "用法: image_operation.sh [保存|加载]"
else
    case $1 in
        保存)
            save
            ;;
        加载)
            load
            ;;
        *)
            echo "无效操作。请输入 '保存' 或 '加载'。"
            ;;
    esac
fi