#!/bin/bash

# Docker镜像源配置脚本
# 描述: 配置Docker使用不同的镜像源，以私有仓库为主，其他源为备

set -e

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

# 镜像源配置 - 以您的私有源为主
PRIVATE_REGISTRY="github.vps7k7k.xyz"
PRIVATE_MIRROR="https://github.vps7k7k.xyz"

# 备用镜像源
ALIYUN_MIRROR="https://registry.cn-hangzhou.aliyuncs.com"
DOCKER_HUB_MIRROR="https://hub-mirror.c.163.com"
USTC_MIRROR="https://docker.mirrors.ustc.edu.cn"
TENCENT_MIRROR="https://mirror.ccs.tencentyun.com"

# 打印彩色信息
print_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

print_choice() {
    echo -e "${PURPLE}[CHOICE]${NC} $1"
}

print_config() {
    echo -e "${CYAN}[CONFIG]${NC} $1"
}

# 检查Docker是否安装
check_docker() {
    if ! command -v docker &> /dev/null; then
        print_error "Docker 未安装，请先安装Docker"
        exit 1
    fi
    print_info "Docker 已安装，版本: $(docker --version | cut -d' ' -f3 | cut -d',' -f1)"
}

# 显示镜像源选择菜单
show_mirror_menu() {
    echo
    print_step "请选择要使用的Docker镜像源配置："
    echo "=========================================="
    print_choice "1. 私有仓库为主 + 多个备用源（推荐）"
    echo "   - 主源: ${PRIVATE_REGISTRY}"
    echo "   - 备源: 阿里云、网易、中科大、腾讯云"
    echo "   - Docker自动选择最快可用源"
    print_choice "2. 仅私有仓库"
    echo "   - 只使用 ${PRIVATE_REGISTRY}"
    echo "   - 需要账号密码认证"
    echo "   - 访问速度最快"
    print_choice "3. 仅公共镜像源"
    echo "   - 不使用私有仓库"
    echo "   - 使用阿里云、网易、中科大、腾讯云"
    print_choice "4. 官方Docker Hub"
    echo "   - 官方源，镜像最全"
    echo "   - 国外源，访问可能较慢"
    print_choice "5. 自定义配置"
    echo "   - 手动选择镜像源组合"
    echo "=========================================="
}

# 获取用户选择的镜像源
get_mirror_choice() {
    while true; do
        read -p "请输入选择 (1-5): " choice
        case $choice in
            1)
                MIRROR_TYPE="private_primary"
                MIRROR_NAME="私有仓库为主 + 多个备用源"
                break
                ;;
            2)
                MIRROR_TYPE="private_only"
                MIRROR_NAME="仅私有仓库"
                break
                ;;
            3)
                MIRROR_TYPE="public_only"
                MIRROR_NAME="仅公共镜像源"
                break
                ;;
            4)
                MIRROR_TYPE="official"
                MIRROR_NAME="官方Docker Hub"
                break
                ;;
            5)
                MIRROR_TYPE="custom"
                MIRROR_NAME="自定义配置"
                break
                ;;
            *)
                print_error "无效选择，请输入 1-5 的数字"
                ;;
        esac
    done
}

# 自定义配置选择
get_custom_config() {
    echo
    print_step "自定义镜像源配置"
    
    MIRRORS=()
    
    # 选择是否包含私有仓库
    read -p "是否包含私有仓库 ${PRIVATE_REGISTRY}？(y/N): " include_private
    if [[ $include_private =~ ^[Yy]$ ]]; then
        MIRRORS+=("$PRIVATE_MIRROR")
        print_info "已添加私有仓库"
    fi
    
    # 选择公共镜像源
    echo
    echo "请选择要添加的公共镜像源（可多选）："
    echo "1) 阿里云镜像"
    echo "2) 网易云镜像" 
    echo "3) 中科大镜像"
    echo "4) 腾讯云镜像"
    echo "5) 官方Docker Hub"
    
    read -p "请输入选择（多个用空格分隔，如：1 2 3）: " public_choices
    
    for choice in $public_choices; do
        case $choice in
            1)
                MIRRORS+=("$ALIYUN_MIRROR")
                print_info "已添加阿里云镜像"
                ;;
            2)
                MIRRORS+=("$DOCKER_HUB_MIRROR")
                print_info "已添加网易云镜像"
                ;;
            3)
                MIRRORS+=("$USTC_MIRROR")
                print_info "已添加中科大镜像"
                ;;
            4)
                MIRRORS+=("$TENCENT_MIRROR")
                print_info "已添加腾讯云镜像"
                ;;
            5)
                MIRRORS+=("https://registry-1.docker.io")
                print_info "已添加官方Docker Hub"
                ;;
        esac
    done
    
    if [ ${#MIRRORS[@]} -eq 0 ]; then
        print_error "至少需要选择一个镜像源"
        exit 1
    fi
    
    MIRROR_NAME="自定义镜像源（${#MIRRORS[@]}个）"
}

# 交互式输入私有仓库账号密码
get_private_credentials() {
    if [[ "$MIRROR_TYPE" != "private_primary" && "$MIRROR_TYPE" != "private_only" ]]; then
        return 0
    fi
    
    print_step "请输入私有仓库的认证信息"
    
    while true; do
        read -p "请输入用户名: " username
        if [ -n "$username" ]; then
            break
        else
            print_error "用户名不能为空，请重新输入"
        fi
    done
    
    while true; do
        read -s -p "请输入密码: " password
        echo
        if [ -n "$password" ]; then
            break
        else
            print_error "密码不能为空，请重新输入"
        fi
    done
    
    # 确认密码
    while true; do
        read -s -p "请再次输入密码确认: " password_confirm
        echo
        if [ "$password" = "$password_confirm" ]; then
            break
        else
            print_error "两次输入的密码不一致，请重新输入"
        fi
    done
}

# 测试私有仓库登录
test_private_login() {
    if [[ "$MIRROR_TYPE" != "private_primary" && "$MIRROR_TYPE" != "private_only" ]]; then
        return 0
    fi
    
    print_step "测试登录私有仓库..."
    
    if echo "$password" | docker login $PRIVATE_REGISTRY -u "$username" --password-stdin; then
        print_info "私有仓库登录测试成功！"
        return 0
    else
        print_error "私有仓库登录测试失败，请检查账号密码和网络连接"
        return 1
    fi
}

# 配置镜像源
configure_mirrors() {
    print_step "配置Docker镜像源: $MIRROR_NAME"
    
    # 创建Docker配置目录
    sudo mkdir -p /etc/docker
    
    # 备份原有配置
    if [ -f /etc/docker/daemon.json ]; then
        sudo cp /etc/docker/daemon.json /etc/docker/daemon.json.backup.$(date +%Y%m%d%H%M%S)
        print_info "已备份原有配置"
    fi
    
    # 根据不同类型生成配置
    case "$MIRROR_TYPE" in
        "private_primary")
            # 私有仓库为主 + 多个备用源
            cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": [
    "$PRIVATE_MIRROR",
    "$ALIYUN_MIRROR",
    "$DOCKER_HUB_MIRROR",
    "$USTC_MIRROR",
    "$TENCENT_MIRROR"
  ],
  "insecure-registries": ["$PRIVATE_REGISTRY"],
  "debug": false,
  "experimental": false
}
EOF
            ;;
        "private_only")
            # 仅私有仓库
            cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": ["$PRIVATE_MIRROR"],
  "insecure-registries": ["$PRIVATE_REGISTRY"],
  "debug": false,
  "experimental": false
}
EOF
            ;;
        "public_only")
            # 仅公共镜像源
            cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": [
    "$ALIYUN_MIRROR",
    "$DOCKER_HUB_MIRROR",
    "$USTC_MIRROR",
    "$TENCENT_MIRROR"
  ],
  "insecure-registries": [],
  "debug": false,
  "experimental": false
}
EOF
            ;;
        "official")
            # 官方Docker Hub
            cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": ["https://registry-1.docker.io"],
  "insecure-registries": [],
  "debug": false,
  "experimental": false
}
EOF
            ;;
        "custom")
            # 自定义配置
            mirrors_json=$(printf '"%s",' "${MIRRORS[@]}")
            mirrors_json="[${mirrors_json%,}]"
            
            insecure_registries="[]"
            if [[ " ${MIRRORS[@]} " =~ " $PRIVATE_MIRROR " ]]; then
                insecure_registries="[\"$PRIVATE_REGISTRY\"]"
            fi
            
            cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": $mirrors_json,
  "insecure-registries": $insecure_registries,
  "debug": false,
  "experimental": false
}
EOF
            ;;
    esac
    
    print_info "Docker镜像源配置完成"
}

# 重启Docker服务
restart_docker() {
    print_step "重启Docker服务..."
    
    if sudo systemctl restart docker; then
        print_info "Docker服务重启成功"
    else
        print_error "Docker服务重启失败"
        exit 1
    fi
    
    # 等待Docker服务完全启动
    sleep 3
}

# 保存私有仓库认证信息
save_private_credentials() {
    if [[ "$MIRROR_TYPE" != "private_primary" && "$MIRROR_TYPE" != "private_only" ]]; then
        return 0
    fi
    
    print_step "是否保存私有仓库认证信息？"
    echo "如果保存，以后拉取私有镜像时不需要重复输入密码"
    read -p "是否保存认证信息？(y/N): " save_choice
    
    if [[ $save_choice =~ ^[Yy]$ ]]; then
        # 使用Docker login保存认证信息
        if echo "$password" | docker login $PRIVATE_REGISTRY -u "$username" --password-stdin; then
            print_info "私有仓库认证信息已保存到 ~/.docker/config.json"
        else
            print_warning "私有仓库认证信息保存失败"
        fi
    else
        print_info "私有仓库认证信息未保存，拉取私有镜像时需要重新登录"
    fi
}

# 验证配置
verify_configuration() {
    print_step "验证配置..."
    
    # 检查Docker服务状态
    if sudo systemctl is-active --quiet docker; then
        print_info "Docker服务运行正常"
    else
        print_error "Docker服务未运行"
        exit 1
    fi
    
    # 检查配置是否生效
    if sudo docker info 2>/dev/null | grep -q "Registry Mirrors"; then
        print_info "镜像源配置已生效"
        echo
        print_config "当前配置的镜像源："
        sudo docker info 2>/dev/null | grep -A 10 "Registry Mirrors"
    else
        print_warning "镜像源配置可能未生效，请手动检查"
    fi
}

# 显示配置摘要
show_summary() {
    echo
    print_info "=== 配置完成摘要 ==="
    echo -e "${GREEN}✓${NC} Docker镜像源配置完成"
    echo -e "${GREEN}✓${NC} 镜像源类型: ${MIRROR_NAME}"
    
    case "$MIRROR_TYPE" in
        "private_primary")
            echo -e "${GREEN}✓${NC} 主镜像源: ${PRIVATE_REGISTRY}"
            echo -e "${GREEN}✓${NC} 备用镜像源: 阿里云、网易、中科大、腾讯云"
            ;;
        "private_only")
            echo -e "${GREEN}✓${NC} 镜像源: ${PRIVATE_REGISTRY}"
            ;;
        "public_only")
            echo -e "${GREEN}✓${NC} 镜像源: 阿里云、网易、中科大、腾讯云"
            ;;
        "official")
            echo -e "${GREEN}✓${NC} 镜像源: 官方Docker Hub"
            ;;
        "custom")
            echo -e "${GREEN}✓${NC} 自定义镜像源: ${#MIRRORS[@]}个"
            for mirror in "${MIRRORS[@]}"; do
                echo -e "  - ${mirror}"
            done
            ;;
    esac
    
    if [[ "$MIRROR_TYPE" == "private_primary" || "$MIRROR_TYPE" == "private_only" ]]; then
        echo -e "${GREEN}✓${NC} 私有仓库用户名: ${username}"
        if [[ $save_choice =~ ^[Yy]$ ]]; then
            echo -e "${GREEN}✓${NC} 认证信息已保存"
        else
            echo -e "${YELLOW}⚠${NC} 认证信息未保存"
        fi
    fi
    
    echo -e "${GREEN}✓${NC} Docker服务已重启生效"
    
    echo
    print_info "使用说明："
    if [[ "$MIRROR_TYPE" == "private_primary" || "$MIRROR_TYPE" == "private_only" ]]; then
        echo "  私有镜像: docker pull ${PRIVATE_REGISTRY}/镜像名称"
        echo "  公共镜像: docker pull 镜像名称"
    else
        echo "  拉取镜像: docker pull 镜像名称"
    fi
    echo "  查看配置: docker info | grep -A 10 'Registry Mirrors'"
}

# 主函数
main() {
    clear
    echo "=========================================="
    echo "    Docker镜像源配置脚本"
    echo "   主镜像源: ${PRIVATE_REGISTRY}"
    echo "=========================================="
    echo
    
    # 检查Docker
    check_docker
    
    # 显示镜像源菜单
    show_mirror_menu
    
    # 获取用户选择
    get_mirror_choice
    
    # 如果是自定义配置，获取详细选择
    if [ "$MIRROR_TYPE" = "custom" ]; then
        get_custom_config
    fi
    
    # 如果是私有仓库配置，获取认证信息
    if [[ "$MIRROR_TYPE" == "private_primary" || "$MIRROR_TYPE" == "private_only" ]]; then
        get_private_credentials
        if ! test_private_login; then
            print_error "配置中止，请解决登录问题后重新运行脚本"
            exit 1
        fi
    fi
    
    # 配置镜像源
    configure_mirrors
    
    # 重启Docker服务
    restart_docker
    
    # 保存私有仓库认证信息
    if [[ "$MIRROR_TYPE" == "private_primary" || "$MIRROR_TYPE" == "private_only" ]]; then
        save_private_credentials
    fi
    
    # 验证配置
    verify_configuration
    
    # 显示摘要
    show_summary
}

# 运行主函数
main "$@"