#!/bin/bash

# Debian 12 Docker 安装脚本 (使用国内源)
# 从国内镜像源安装 Docker CE
# 支持多源切换和自动选择最佳源

set -e

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

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

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

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

# 定义多个 Docker 源
declare -A DOCKER_SOURCES=(
    ["aliyun"]="https://mirrors.aliyun.com/docker-ce"
    ["tsinghua"]="https://mirrors.tuna.tsinghua.edu.cn/docker-ce"
    ["ustc"]="https://mirrors.ustc.edu.cn/docker-ce"
    ["huawei"]="https://mirrors.huaweicloud.com/docker-ce"
    ["163"]="https://mirrors.163.com/docker-ce"
    ["azure"]="https://mirror.azure.cn/docker-ce"
    ["official"]="https://download.docker.com"
)

# 源优先级顺序
DOCKER_SOURCE_ORDER=("aliyun" "ustc" "huawei" "163" "tsinghua" "azure" "official")

# 当前使用的源
CURRENT_SOURCE=""
CURRENT_SOURCE_URL=""

# 测试单个源的连接性
test_source() {
    local source_name="$1"
    local source_url="${DOCKER_SOURCES[$source_name]}"
    local codename=$(lsb_release -cs)
    local arch=$(dpkg --print-architecture)
    
    local test_url="${source_url}/linux/debian/dists/${codename}/stable/binary-${arch}/Packages"
    
    # 测试连接
    local start_time=$(date +%s%N)
    if curl -sfL --connect-timeout 5 --max-time 15 "$test_url" -o /dev/null 2>/dev/null; then
        local end_time=$(date +%s%N)
        local time_diff=$(( (end_time - start_time) / 1000000 ))
        echo "$time_diff"
        return 0
    fi
    return 1
}

# 测试所有源并选择最佳
select_best_source() {
    log_info "测试 Docker 源连接速度..."
    
    local best_source=""
    local best_time=999999
    
    for source_name in "${DOCKER_SOURCE_ORDER[@]}"; do
        local source_url="${DOCKER_SOURCES[$source_name]}"
        printf "  %-12s %-50s " "$source_name" "$source_url"
        
        if time_ms=$(test_source "$source_name"); then
            echo -e "${GREEN}✓${NC} ${time_ms}ms"
            if (( time_ms < best_time )); then
                best_time=$time_ms
                best_source=$source_name
            fi
        else
            echo -e "${RED}✗${NC} 连接失败"
        fi
    done
    
    if [[ -n "$best_source" ]]; then
        log_success "选择最佳源: $best_source (${best_time}ms)"
        CURRENT_SOURCE="$best_source"
        CURRENT_SOURCE_URL="${DOCKER_SOURCES[$best_source]}"
        return 0
    else
        log_error "所有源均不可用"
        return 1
    fi
}

# 使用指定源
use_source() {
    local source_name="$1"
    
    if [[ -z "${DOCKER_SOURCES[$source_name]}" ]]; then
        log_error "未知的源: $source_name"
        log_info "可用源: ${!DOCKER_SOURCES[*]}"
        return 1
    fi
    
    log_info "测试指定源: $source_name"
    if test_source "$source_name" >/dev/null; then
        CURRENT_SOURCE="$source_name"
        CURRENT_SOURCE_URL="${DOCKER_SOURCES[$source_name]}"
        log_success "使用源: $source_name ($CURRENT_SOURCE_URL)"
        return 0
    else
        log_error "源 $source_name 不可用"
        return 1
    fi
}

# 尝试下一个源
try_next_source() {
    local current="$CURRENT_SOURCE"
    local found_current=false
    
    for source_name in "${DOCKER_SOURCE_ORDER[@]}"; do
        if [[ "$found_current" == "true" ]]; then
            if test_source "$source_name" >/dev/null; then
                CURRENT_SOURCE="$source_name"
                CURRENT_SOURCE_URL="${DOCKER_SOURCES[$source_name]}"
                log_info "切换到备用源: $source_name"
                return 0
            fi
        fi
        if [[ "$source_name" == "$current" ]]; then
            found_current=true
        fi
    done
    
    # 如果没有找到，从头开始尝试
    for source_name in "${DOCKER_SOURCE_ORDER[@]}"; do
        if [[ "$source_name" != "$current" ]]; then
            if test_source "$source_name" >/dev/null; then
                CURRENT_SOURCE="$source_name"
                CURRENT_SOURCE_URL="${DOCKER_SOURCES[$source_name]}"
                log_info "切换到备用源: $source_name"
                return 0
            fi
        fi
    done
    
    return 1
}

# 检查系统要求
check_system() {
    log_info "检查系统要求..."

    # 检查是否为 Debian 系统
    if ! command -v lsb_release &> /dev/null; then
        apt update && apt install -y lsb-release
    fi

    local distro=$(lsb_release -si)
    local version=$(lsb_release -sr)

    if [[ "$distro" != "Debian" ]]; then
        log_error "此脚本仅支持 Debian 系统"
        exit 1
    fi

    if [[ "${version%%.*}" -lt 10 ]]; then
        log_error "此脚本需要 Debian 10 或更高版本"
        exit 1
    fi

    # 检查架构
    local arch=$(dpkg --print-architecture)
    if [[ "$arch" != "amd64" && "$arch" != "arm64" ]]; then
        log_error "不支持的架构: $arch"
        exit 1
    fi

    log_success "系统检查通过: $distro $version ($arch)"
}

# 卸载旧版本
uninstall_old_versions() {
    log_info "卸载可能存在的旧版本 Docker..."

    # 停止所有运行中的容器
    if command -v docker &> /dev/null; then
        docker stop $(docker ps -aq) 2>/dev/null || true
    fi

    # 卸载旧版本
    apt remove -y docker docker-engine docker.io containerd runc docker-compose 2>/dev/null || true

    # 清理相关文件
    rm -rf /var/lib/docker /etc/docker
    rm -f /etc/apt/sources.list.d/docker.list
    rm -f /etc/apt/keyrings/docker.gpg

    log_success "旧版本清理完成"
}

# 安装依赖包
install_dependencies() {
    log_info "安装依赖包..."

    apt update
    apt install -y \
        ca-certificates \
        curl \
        gnupg \
        lsb-release \
        apt-transport-https \
        software-properties-common

    log_success "依赖包安装完成"
}

# 添加 Docker 仓库
add_docker_repository() {
    if [[ -z "$CURRENT_SOURCE_URL" ]]; then
        log_error "未选择 Docker 源"
        return 1
    fi
    
    log_info "添加 Docker 仓库 (使用 $CURRENT_SOURCE 源)..."
    log_info "源地址: $CURRENT_SOURCE_URL"

    # 创建密钥目录
    mkdir -p /etc/apt/keyrings
    
    # 移除旧的密钥和仓库配置
    rm -f /etc/apt/keyrings/docker.gpg
    rm -f /etc/apt/sources.list.d/docker.list

    # 下载并添加 GPG 密钥
    log_info "下载 GPG 密钥..."
    if ! curl -fsSL "${CURRENT_SOURCE_URL}/linux/debian/gpg" | gpg --dearmor -o /etc/apt/keyrings/docker.gpg 2>/dev/null; then
        log_warning "从 $CURRENT_SOURCE 下载 GPG 密钥失败，尝试官方源..."
        curl -fsSL "https://download.docker.com/linux/debian/gpg" | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
    fi
    
    chmod a+r /etc/apt/keyrings/docker.gpg

    # 添加仓库
    local arch=$(dpkg --print-architecture)
    local codename=$(lsb_release -cs)
    
    echo "deb [arch=${arch} signed-by=/etc/apt/keyrings/docker.gpg] ${CURRENT_SOURCE_URL}/linux/debian ${codename} stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null

    log_success "Docker 仓库添加完成 (源: $CURRENT_SOURCE)"
}

# 安装 Docker
install_docker() {
    log_info "安装 Docker CE..."
    
    local max_retries=3
    local retry_count=0
    
    while (( retry_count < max_retries )); do
        # 更新包索引
        log_info "更新包索引..."
        if ! apt update 2>&1; then
            log_warning "apt update 失败"
        fi

        # 尝试安装 Docker CE
        log_info "尝试从 $CURRENT_SOURCE 安装 Docker..."
        
        # 使用 set +e 暂时禁用错误退出
        set +e
        apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin 2>&1
        local install_result=$?
        set -e
        
        if [[ $install_result -eq 0 ]]; then
            log_success "Docker CE 安装完成 (源: $CURRENT_SOURCE)"
            return 0
        fi
        
        retry_count=$((retry_count + 1))
        log_warning "安装失败 (尝试 $retry_count/$max_retries)"
        
        if (( retry_count < max_retries )); then
            log_info "尝试切换到其他源..."
            
            if try_next_source; then
                # 重新配置仓库
                add_docker_repository
            else
                log_error "没有更多可用的源"
                break
            fi
        fi
    done
    
    log_error "Docker 安装失败，已尝试所有可用源"
    log_info "建议手动检查网络连接或稍后重试"
    exit 1
}

# 配置 Docker
configure_docker() {
    log_info "配置 Docker..."

    # 创建配置目录
    mkdir -p /etc/docker

    # 配置 daemon.json
    cat > /etc/docker/daemon.json << 'EOF'
{
  "registry-mirrors": [
    "https://docker-proxy.syy.1msoft.cn",
    "https://registry.docker-cn.com",
    "https://docker.mirrors.ustc.edu.cn",
    "https://hub-mirror.c.163.com",
    "https://mirror.baidubce.com"
  ],
  "log-driver": "json-file",
  "log-opts": {
    "max-size": "100m",
    "max-file": "3"
  },
  "storage-driver": "overlay2",
  "exec-opts": ["native.cgroupdriver=systemd"],
  "iptables": false
}
EOF

    # 配置 containerd
    mkdir -p /etc/containerd
    containerd config default | tee /etc/containerd/config.toml > /dev/null

    # 重新加载配置
    systemctl daemon-reload

    log_success "Docker 配置完成"
}

# 启动服务
start_services() {
    log_info "启动 Docker 服务..."

    # 启动服务
    systemctl enable docker
    systemctl enable containerd
    systemctl start docker
    systemctl start containerd

    # 等待服务启动
    sleep 2

    # 检查服务状态
    if systemctl is-active --quiet docker; then
        log_success "Docker 服务启动成功"
    else
        log_error "Docker 服务启动失败"
        journalctl -u docker --no-pager -n 20
        exit 1
    fi
}

# 配置用户权限
configure_user_permissions() {
    log_info "配置用户权限..."

    # 获取当前用户
    local current_user=""
    if [[ -n "$SUDO_USER" ]]; then
        current_user="$SUDO_USER"
    elif [[ -n "$USER" && "$USER" != "root" ]]; then
        current_user="$USER"
    fi

    if [[ -n "$current_user" ]]; then
        # 检查用户是否存在
        if id "$current_user" &>/dev/null; then
            usermod -aG docker "$current_user"
            log_success "已将用户 $current_user 添加到 docker 组"
            log_warning "请重新登录或运行 'newgrp docker' 以使权限生效"
        fi
    fi
}

# 测试安装
test_installation() {
    log_info "测试 Docker 安装..."

    # 检查版本
    local docker_version=$(docker --version)
    local docker_compose_version=$(docker compose version 2>/dev/null || echo "docker-compose plugin")

    log_success "Docker 版本: $docker_version"
    log_info "Docker Compose: $docker_compose_version"

    # 运行测试容器
    log_info "运行测试容器..."
    if docker run --rm hello-world >/dev/null 2>&1; then
        log_success "Docker 测试通过！"
    else
        log_warning "Docker 运行测试失败，但安装可能成功。请检查网络连接。"
    fi

    # 显示使用信息
    cat << 'EOF'

Docker 安装完成！

常用命令:
  docker --version              # 查看版本
  docker run hello-world        # 运行测试容器
  docker ps -a                  # 查看所有容器
  docker images                 # 查看镜像
  docker pull <image>           # 拉取镜像
  docker build -t <name> .      # 构建镜像

服务管理:
  sudo systemctl start docker   # 启动服务
  sudo systemctl stop docker    # 停止服务
  sudo systemctl restart docker # 重启服务

配置说明:
  - 已配置国内镜像加速器
  - 日志轮转已配置 (最大100M，保留3个文件)
  - 使用 overlay2 存储驱动
  - 已启用 systemd cgroup 驱动

EOF
}

# 列出可用源
list_sources() {
    echo "可用的 Docker 源:"
    echo "=========================================="
    for source_name in "${DOCKER_SOURCE_ORDER[@]}"; do
        printf "  %-12s %s\n" "$source_name" "${DOCKER_SOURCES[$source_name]}"
    done
    echo ""
    echo "默认优先级: ${DOCKER_SOURCE_ORDER[*]}"
}

# 显示帮助信息
show_help() {
    cat << EOF
Debian 12 Docker 安装脚本

此脚本支持多个国内镜像源，自动选择最佳源或手动指定源。

用法: $0 [选项]

选项:
    -h, --help          显示此帮助信息
    -f, --force         强制重新安装（不检查现有安装）
    -s, --source NAME   指定使用的源 (aliyun, tsinghua, ustc, huawei, 163, azure, official)
    -t, --test          仅测试所有源的连接速度
    -l, --list          列出所有可用源

可用源:
    aliyun      阿里云镜像 (推荐)
    tsinghua    清华大学镜像
    ustc        中科大镜像
    huawei      华为云镜像
    163         网易镜像
    azure       微软 Azure 镜像
    official    Docker 官方源

示例:
    $0                      # 自动选择最佳源安装
    $0 -s aliyun            # 使用阿里云源安装
    $0 -s ustc --force      # 使用中科大源强制重装
    $0 --test               # 测试所有源速度
    $0 --list               # 列出所有可用源

EOF
}

# 主函数
main() {
    local force_install=false
    local specified_source=""
    local test_only=false

    # 解析参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -f|--force)
                force_install=true
                shift
                ;;
            -s|--source)
                specified_source="$2"
                shift 2
                ;;
            -t|--test)
                test_only=true
                shift
                ;;
            -l|--list)
                list_sources
                exit 0
                ;;
            *)
                log_error "未知选项: $1"
                show_help
                exit 1
                ;;
        esac
    done

    # 检查是否为 root 用户
    if [[ $EUID -ne 0 ]]; then
        log_error "请使用 root 用户或 sudo 运行此脚本"
        exit 1
    fi

    # 仅测试模式
    if [[ "$test_only" == true ]]; then
        select_best_source
        exit 0
    fi

    # 检查是否已安装
    if command -v docker &> /dev/null && [[ "$force_install" != true ]]; then
        log_warning "Docker 已经安装。如需重新安装，请使用 --force 选项。"
        docker --version
        exit 0
    fi

    log_info "开始安装 Docker CE..."
    echo

    check_system
    
    # 选择源
    if [[ -n "$specified_source" ]]; then
        if ! use_source "$specified_source"; then
            log_warning "指定源不可用，自动选择最佳源..."
            select_best_source || exit 1
        fi
    else
        select_best_source || exit 1
    fi
    
    echo
    
    uninstall_old_versions
    install_dependencies
    add_docker_repository
    install_docker
    configure_docker
    start_services
    configure_user_permissions
    test_installation

    echo
    log_success "Docker 安装完成！(使用源: $CURRENT_SOURCE)"
}

# 执行主函数
main "$@"