#!/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"
ALIYUN_MIRROR="https://registry.cn-hangzhou.aliyuncs.com"
DOCKER_HUB_MIRROR="https://hub-mirror.c.163.com"
OFFICIAL_REGISTRY="https://registry-1.docker.io"

# 打印彩色信息
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. 私有仓库 (${PRIVATE_REGISTRY})"
    echo "   - 需要账号密码认证"
    echo "   - 访问速度最快"
    echo "   - 包含定制镜像"
    print_choice "2. 阿里云镜像"
    echo "   - 国内访问速度快"
    echo "   - 无需认证"
    echo "   - 镜像同步及时"
    print_choice "3. 网易云镜像"
    echo "   - 国内镜像源"
    echo "   - 无需认证"
    echo "   - 稳定性好"
    print_choice "4. 官方Docker Hub"
    echo "   - 官方源，镜像最全"
    echo "   - 国外源，访问可能较慢"
    print_choice "5. 多个镜像源（推荐）"
    echo "   - 同时配置多个镜像源"
    echo "   - Docker会自动选择最快的源"
    echo "=========================================="
}

# 获取用户选择的镜像源
get_mirror_choice() {
    while true; do
        read -p "请输入选择 (1-5): " choice
        case $choice in
            1)
                MIRROR_TYPE="private"
                MIRROR_NAME="私有仓库"
                MIRROR_URL="https://${PRIVATE_REGISTRY}"
                break
                ;;
            2)
                MIRROR_TYPE="aliyun"
                MIRROR_NAME="阿里云镜像"
                MIRROR_URL=$ALIYUN_MIRROR
                break
                ;;
            3)
                MIRROR_TYPE="netease"
                MIRROR_NAME="网易云镜像"
                MIRROR_URL=$DOCKER_HUB_MIRROR
                break
                ;;
            4)
                MIRROR_TYPE="official"
                MIRROR_NAME="官方Docker Hub"
                MIRROR_URL=$OFFICIAL_REGISTRY
                break
                ;;
            5)
                MIRROR_TYPE="multiple"
                MIRROR_NAME="多个镜像源"
                break
                ;;
            *)
                print_error "无效选择，请输入 1-5 的数字"
                ;;
        esac
    done
}

# 交互式输入私有仓库账号密码
get_private_credentials() {
    if [ "$MIRROR_TYPE" != "private" ]; 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" ]; 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_single_mirror() {
    local mirror_url=$1
    
    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
    
    # 创建新的配置
    if [ "$MIRROR_TYPE" = "private" ]; then
        cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": ["$mirror_url"],
  "insecure-registries": ["$PRIVATE_REGISTRY"],
  "debug": false,
  "experimental": false
}
EOF
    else
        cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": ["$mirror_url"],
  "insecure-registries": [],
  "debug": false,
  "experimental": false
}
EOF
    fi
    
    print_info "Docker镜像源配置完成"
}

# 配置多个镜像源
configure_multiple_mirrors() {
    print_step "配置多个Docker镜像源"
    
    # 创建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
    
    # 创建多个镜像源配置
    cat << EOF | sudo tee /etc/docker/daemon.json > /dev/null
{
  "registry-mirrors": [
    "$ALIYUN_MIRROR",
    "$DOCKER_HUB_MIRROR",
    "https://docker.mirrors.ustc.edu.cn",
    "https://mirror.ccs.tencentyun.com"
  ],
  "insecure-registries": ["$PRIVATE_REGISTRY"],
  "debug": false,
  "experimental": false
}
EOF
    
    print_info "多个Docker镜像源配置完成"
    print_info "已配置的镜像源："
    print_config "  - 阿里云镜像"
    print_config "  - 网易云镜像"
    print_config "  - 中科大镜像"
    print_config "  - 腾讯云镜像"
}

# 重启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" ]; 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
    
    # 测试拉取镜像（可选）
    if [ "$MIRROR_TYPE" = "private" ]; then
        print_step "是否测试拉取私有仓库的nginx镜像？"
        read -p "进行测试？(y/N): " test_choice
        
        if [[ $test_choice =~ ^[Yy]$ ]]; then
            print_info "开始拉取测试镜像..."
            if sudo docker pull ${PRIVATE_REGISTRY}/nginx; then
                print_info "私有仓库镜像拉取测试成功！"
            else
                print_warning "私有仓库镜像拉取测试失败，请检查网络连接和认证信息"
            fi
        fi
    else
        print_step "是否测试拉取公共镜像？"
        read -p "进行测试？(y/N): " test_choice
        
        if [[ $test_choice =~ ^[Yy]$ ]]; then
            print_info "开始拉取测试镜像..."
            if sudo docker pull hello-world; then
                print_info "公共镜像拉取测试成功！"
            else
                print_warning "公共镜像拉取测试失败，请检查网络连接"
            fi
        fi
    fi
}

# 显示配置摘要
show_summary() {
    echo
    print_info "=== 配置完成摘要 ==="
    echo -e "${GREEN}✓${NC} Docker镜像源配置完成"
    echo -e "${GREEN}✓${NC} 镜像源类型: ${MIRROR_NAME}"
    
    if [ "$MIRROR_TYPE" = "private" ]; then
        echo -e "${GREEN}✓${NC} 私有仓库地址: ${PRIVATE_REGISTRY}"
        echo -e "${GREEN}✓${NC} 用户名: ${username}"
        if [[ $save_choice =~ ^[Yy]$ ]]; then
            echo -e "${GREEN}✓${NC} 认证信息已保存，后续使用无需重复登录"
        else
            echo -e "${YELLOW}⚠${NC} 认证信息未保存，拉取私有镜像时需要重新登录"
        fi
    elif [ "$MIRROR_TYPE" = "multiple" ]; then
        echo -e "${GREEN}✓${NC} 已配置多个镜像源"
        echo -e "${GREEN}✓${NC} Docker会自动选择最快的镜像源"
    else
        echo -e "${GREEN}✓${NC} 镜像源地址: ${MIRROR_URL}"
    fi
    
    echo -e "${GREEN}✓${NC} Docker服务已重启生效"
    
    echo
    print_info "使用说明："
    if [ "$MIRROR_TYPE" = "private" ]; then
        echo "  私有镜像: docker pull ${PRIVATE_REGISTRY}/镜像名称"
        echo "  公共镜像: docker pull 镜像名称"
    else
        echo "  拉取镜像: docker pull 镜像名称"
    fi
}

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

# 运行主函数
main "$@"