#!/bin/bash
# pull-docker-image.sh
# 主要功能：
#   - 根据用户参数拉取指定 Docker 镜像，生成 tar 包，并支持 gzip 或 zstd 压缩。
#
# 具体参数与用法请见：pull-docker-image.sh -h

set -euo pipefail

# Default values.
IMAGE=""
DEST="."
NEWNAME=""
USE_GZIP=false
USE_ZSTD=false
FORCE=false
PROXY=""

usage() {
    cat <<EOF
Usage:  $(basename "$0") -t image:tag [-d destination] [-n newname:tag] [-g] [-z] 
                                      [-f] [-p proxy]
  -t    (Required) Docker image name and tag to pull.
  -d    Destination of pulled image file。
          - If a directory is given, a new file named as image name and tag 
            (with ':' replaced by '_') under the given directory is created.
            Suffixes are inferred from compression options if specified.
          - If a file path is given, the given file path is used directly.
          - To forcefully overwrite existing file, use '-f' option.
          - Default: . (Current directory).
  -n    Image name and tag in output image file. Keep the same as '-t' option
        if not specified.
  -g    Compress image file with gzip.
  -z    Compress image file with zstd.
  -f    Forcefully overwrite existing destination file。
  -p    Set up network proxies (through 'HTTP_PROXY' and 'HTTPS_PROXY' variables).
  -h    Show this help information.
  
Example:
  1. Pulling 'nginx:mainline' image to existed directory '~/nginx-images/'.
         $(basename "$0") -t nginx:mainline -d ~/nginx-images/
EOF
}

# 解析命令行选项
parse_args() {
    while getopts "t:d:n:gzfp:h" opt; do
        case "$opt" in
            t)
                IMAGE="$OPTARG"
                ;;
            d)
                DEST="$OPTARG"
                ;;
            n)
                NEWNAME="$OPTARG"
                ;;
            g)
                USE_GZIP=true
                ;;
            z)
                USE_ZSTD=true
                ;;
            f)
                FORCE=true
                ;;
            p)
                PROXY="$OPTARG"
                ;;
            h)
                usage
                exit 0
                ;;
            *)
                usage
                exit 1
                ;;
        esac
    done

    if [[ -z "$IMAGE" ]]; then
        echo "Error: '-t' option (image name and tag to pull) must be specified." >&2
        usage
        exit 1
    fi

    # 如果未指定 -n，则使用 -t 的值
    if [[ -z "$NEWNAME" ]]; then
        NEWNAME="$IMAGE"
    fi

    # 压缩方法 -g 和 -z 不能同时指定
    if $USE_GZIP && $USE_ZSTD; then
        echo "Error: '-g' and '-z' cannot be used together." >&2
        exit 1
    fi
}

# 检查是否安装了 skopeo
check_skopeo() {
    if ! command -v skopeo >/dev/null 2>&1; then
        echo "Error: skopeo is not installed, or not in PATH." >&2
        exit 1
    fi
}

# 设置网络代理
setup_proxy() {
    if [[ -n "$PROXY" ]]; then
        export HTTP_PROXY="$PROXY"
        export HTTPS_PROXY="$PROXY"
        echo "Use proxy at $PROXY"
    fi
}

# 根据目标参数确定最终的目标文件路径
determine_target_file() {
    local target=""
    if [[ -d "$DEST" ]]; then
        local base_name
        base_name=$(echo "$IMAGE" | tr ':' '_')
        target="${DEST%/}/$base_name.tar"
    else
        target="$DEST"
    fi
    echo "$target"
}

# 检查目标文件是否存在冲突
check_conflict() {
    local file="$1"
    if [[ -f "$file" ]]; then
        if ! $FORCE; then
            echo "Error: destination file '$file' exists, use '-f' option to overwrite." >&2
            exit 1
        else
            echo "Warning: destination file '$file' exists, override as '-f' specified."
        fi
    fi
}

# 使用 skopeo 拉取 Docker 镜像
pull_image() {
    local target_file="$1"
    echo "Start pulling $IMAGE to $target_file with name $NEWNAME in archive..."
    if ! skopeo copy "docker://$IMAGE" "docker-archive:$target_file:$NEWNAME"; then
        echo "Error: skopeo failed during pulling image" >&2
        exit 1
    fi
    echo "Image is saved to $target_file successfully."
}

# 根据压缩选项对镜像文件进行压缩
compress_image() {
    local target_file="$1"
    local final_file=""
    if $USE_GZIP || $USE_ZSTD; then
        if $USE_GZIP; then
            local compressed="${target_file}.gz"
            check_conflict "$compressed"
            echo "Compress $target_file with gzip..."
            if ! gzip -f "$target_file"; then
                echo "Error: failed to compress image with gzip." >&2
                exit 1
            fi
            final_file="$compressed"
        elif $USE_ZSTD; then
            local compressed="${target_file}.zst"
            check_conflict "$compressed"
            echo "Compress $target_file with zstd..."
            if ! zstd -q -o "$compressed" "$target_file"; then
                echo "Error: failed to compress image with zstd." >&2
                exit 1
            fi
            rm -f "$target_file"
            final_file="$compressed"
        fi
        echo "Compressed image is saved to $final_file."
    else
        final_file="$target_file"
    fi
}

# 主函数
main() {
    parse_args "$@"
    check_skopeo
    setup_proxy

    local target_file
    target_file=$(determine_target_file)
    check_conflict "$target_file"

    # 同时检查压缩后的文件是否存在冲突（如果有压缩选项）
    if $USE_GZIP; then
        check_conflict "${target_file}.gz"
    elif $USE_ZSTD; then
        check_conflict "${target_file}.zst"
    fi

    pull_image "$target_file"
    compress_image "$target_file"
    
    echo "OK."
}

# 主入口
main "$@"