#!/bin/bash

# dockerpull.sh - Docker镜像拉取脚本
# 支持从 docker.890805.xyz 拉取镜像并重命名为原始名称

SCRIPT_NAME="dockerpull"
INSTALL_PATH="/usr/local/bin/$SCRIPT_NAME"
REGISTRY="docker.1ms.run"

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

# 显示帮助信息
show_help() {
    cat << EOF
用法: $0 [选项] <镜像名称1> [镜像名称2] ...

选项:
    install     安装脚本到系统路径 ($INSTALL_PATH)
    uninstall   从系统路径卸载脚本
    -h, --help  显示此帮助信息

示例:
    $0 hello-world                    # 拉取 library/hello-world
    $0 nginx alpine                   # 拉取多个镜像
    $0 easysoft/zentao mysql:5.7      # 拉取指定版本
    $0 install                        # 安装到系统
    $0 uninstall                      # 从系统卸载

说明:
    - 镜像会从 $REGISTRY 拉取
    - 拉取后会重命名为原始镜像名
    - 如果镜像名不包含'/'，会自动添加'library/'前缀
EOF
}

# 安装脚本
install_script() {
    if [ "$(id -u)" -ne 0 ]; then
        echo -e "${RED}错误: 安装需要root权限，请使用sudo运行${NC}"
        exit 1
    fi
    
    if [ ! -f "$0" ]; then
        echo -e "${RED}错误: 无法找到当前脚本文件${NC}"
        exit 1
    fi
    
    cp "$0" "$INSTALL_PATH"
    chmod +x "$INSTALL_PATH"
    
    if [ -f "$INSTALL_PATH" ]; then
        echo -e "${GREEN}✓ 脚本已成功安装到 $INSTALL_PATH${NC}"
        echo -e "${BLUE}现在你可以在任何地方使用 '$SCRIPT_NAME' 命令${NC}"
    else
        echo -e "${RED}✗ 安装失败${NC}"
        exit 1
    fi
}

# 卸载脚本
uninstall_script() {
    if [ "$(id -u)" -ne 0 ]; then
        echo -e "${RED}错误: 卸载需要root权限，请使用sudo运行${NC}"
        exit 1
    fi
    
    if [ -f "$INSTALL_PATH" ]; then
        rm "$INSTALL_PATH"
        echo -e "${GREEN}✓ 脚本已成功从 $INSTALL_PATH 卸载${NC}"
    else
        echo -e "${YELLOW}⚠ 脚本未安装或已被删除${NC}"
    fi
}

# 检查Docker是否安装
check_docker() {
    if ! command -v docker >/dev/null 2>&1; then
        echo -e "${RED}错误: Docker未安装或未在PATH中找到${NC}"
        exit 1
    fi
    
    if ! docker info >/dev/null 2>&1; then
        echo -e "${RED}错误: Docker守护进程未运行或权限不足${NC}"
        echo -e "${YELLOW}提示: 请启动Docker或将当前用户添加到docker组${NC}"
        exit 1
    fi
}

# 格式化镜像名称
format_image_name() {
    local image="$1"
    
    # 如果镜像名不包含斜杠，添加library前缀
    case "$image" in
        */*) echo "$image" ;;
        *) echo "library/$image" ;;
    esac
}

# 显示进度条
show_progress() {
    current=$1
    total=$2
    image_name="$3"
    width=40
    
    # 防止除零错误
    if [ $total -eq 0 ]; then
        percentage=0
        filled=0
    else
        percentage=$((current * 100 / total))
        filled=$((percentage * width / 100))
    fi
    
    printf "\r${BLUE}整体进度 [${current}/${total}] ${NC}"
    printf "["
    progress_i=0
    while [ $progress_i -lt $filled ]; do
        printf "█"
        progress_i=$((progress_i + 1))
    done
    while [ $progress_i -lt $width ]; do
        printf "░"
        progress_i=$((progress_i + 1))
    done
    printf "] %d%% " "$percentage"
    
    if [ -n "$image_name" ]; then
        printf "当前: %s" "$image_name"
    fi
}

# 拉取单个镜像
pull_image() {
    local original_image="$1"
    local formatted_image=$(format_image_name "$original_image")
    local proxy_image="$REGISTRY/$formatted_image"
    
    echo -e "\n${YELLOW}📦 拉取镜像: $original_image${NC}"
    echo -e "${BLUE}   从源: $proxy_image${NC}"
    echo -e "${BLUE}   Docker拉取进度:${NC}"
    
    # 拉取镜像，显示详细进度
    if docker pull "$proxy_image"; then
        echo -e "${GREEN}✓ 拉取成功${NC}"
        
        # 重命名镜像
        echo -e "${BLUE}🏷️  重命名为: $original_image${NC}"
        if docker tag "$proxy_image" "$original_image"; then
            echo -e "${GREEN}✓ 重命名成功${NC}"
            
            # 删除代理镜像
            echo -e "${BLUE}🗑️  清理代理镜像...${NC}"
            if docker rmi "$proxy_image" >/dev/null 2>&1; then
                echo -e "${GREEN}✓ 代理镜像已清理${NC}"
            fi
            echo -e "${GREEN}✓ $original_image 处理完成${NC}"
            return 0
        else
            echo -e "${RED}✗ 重命名失败${NC}"
            return 1
        fi
    else
        echo -e "${RED}✗ 拉取失败: $original_image${NC}"
        return 1
    fi
}

# 主函数
main() {
    # 处理特殊参数
    case "$1" in
        "install")
            install_script
            exit 0
            ;;
        "uninstall")
            uninstall_script
            exit 0
            ;;
        "-h"|"--help"|"help")
            show_help
            exit 0
            ;;
        "")
            echo -e "${RED}错误: 请提供要拉取的镜像名称${NC}"
            echo -e "${BLUE}使用 '$0 --help' 查看帮助信息${NC}"
            exit 1
            ;;
    esac
    
    # 检查Docker环境
    check_docker
    
    # 获取镜像列表（直接使用位置参数）
    total_images=$#
    successful_pulls=0
    failed_pulls=0
    
    echo -e "${BLUE}=================== Docker镜像拉取工具 ===================${NC}"
    echo -e "${BLUE}镜像源: $REGISTRY${NC}"
    echo -e "${BLUE}待拉取镜像数量: $total_images${NC}"
    echo -e "${BLUE}======================================================${NC}"
    
    # 显示初始进度
    show_progress 0 $total_images
    
    # 拉取每个镜像
    image_count=1
    for arg in "$@"; do
        current_image="$arg"
        echo -e "\n${BLUE}=================== [$image_count/$total_images] ===================${NC}"
        show_progress $((image_count - 1)) $total_images "$current_image"
        
        if pull_image "$current_image"; then
            successful_pulls=$((successful_pulls + 1))
        else
            failed_pulls=$((failed_pulls + 1))
        fi
        
        # 显示当前完成的进度
        echo -e "${BLUE}============== 镜像 $image_count/$total_images 处理完毕 ==============${NC}"
        image_count=$((image_count + 1))
    done
    
    # 显示最终进度
    show_progress $total_images $total_images "完成"
    echo -e "\n"
    
    # 显示汇总结果
    echo -e "${BLUE}=================== 拉取结果汇总 ===================${NC}"
    echo -e "${GREEN}✓ 成功拉取: $successful_pulls 个镜像${NC}"
    
    if [ $failed_pulls -gt 0 ]; then
        echo -e "${RED}✗ 拉取失败: $failed_pulls 个镜像${NC}"
    fi
    
    echo -e "${BLUE}=================================================${NC}"
    
    # 设置退出码
    if [ $failed_pulls -gt 0 ]; then
        exit 1
    else
        echo -e "${GREEN}🎉 所有镜像拉取完成！${NC}"
        exit 0
    fi
}

# 运行主函数
main "$@"