#!/bin/bash

# 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'

# 日志函数
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"
}

log_title() {
    echo -e "${PURPLE}[DOCKER MIRRORS]${NC} $1"
}

# 镜像源配置
MIRRORS_aliyun="https://registry.cn-hangzhou.aliyuncs.com"
MIRRORS_tencent="https://mirror.ccs.tencentyun.com"
MIRRORS_huawei="https://mirror-1.trustie.net"
MIRRORS_netease="https://hub-mirror.c.163.com"
MIRRORS_ustc="https://docker.mirrors.ustc.edu.cn"
MIRRORS_docker_cn="https://registry.docker-cn.com"
MIRRORS_baidu="https://mirror.baidubce.com"
MIRRORS_tsinghua="https://mirrors.tuna.tsinghua.edu.cn"

# 镜像源名称数组
MIRROR_NAMES=("aliyun" "tencent" "huawei" "netease" "ustc" "docker_cn" "baidu" "tsinghua")

# 配置文件路径
DOCKER_CONFIG="/etc/docker/daemon.json"
DOCKER_CONFIG_BACKUP="/etc/docker/daemon.json.backup.$(date +%Y%m%d_%H%M%S)"

# 检查 Docker 是否安装
check_docker() {
    if ! command -v docker &> /dev/null; then
        log_error "Docker 未安装，请先安装 Docker"
        echo "运行以下命令安装 Docker:"
        echo "  sudo ./scripts/packages/install_docker.sh"
        exit 1
    fi

    if ! systemctl is-active --quiet docker; then
        log_warning "Docker 服务未运行，正在启动..."
        systemctl start docker
        sleep 2
        if ! systemctl is-active --quiet docker; then
            log_error "无法启动 Docker 服务"
            exit 1
        fi
        log_success "Docker 服务已启动"
    fi
}

# 检查 Docker 配置文件
check_config() {
    if [[ ! -f "$DOCKER_CONFIG" ]]; then
        log_info "创建 Docker 配置文件..."
        mkdir -p /etc/docker
        echo '{}' > "$DOCKER_CONFIG"
    fi
}

# 备份当前配置
backup_config() {
    if [[ -f "$DOCKER_CONFIG" ]]; then
        cp "$DOCKER_CONFIG" "$DOCKER_CONFIG_BACKUP"
        log_info "配置已备份到: $DOCKER_CONFIG_BACKUP"
    fi
}

# 获取当前配置
get_current_config() {
    if [[ -f "$DOCKER_CONFIG" ]]; then
        cat "$DOCKER_CONFIG"
    else
        echo "{}"
    fi
}

# 更新配置文件
update_config() {
    local new_config="$1"
    echo "$new_config" > "$DOCKER_CONFIG"
    log_info "Docker 配置已更新"
}

# 重启 Docker 服务
restart_docker() {
    log_info "重启 Docker 服务..."
    systemctl daemon-reload
    systemctl restart docker

    # 等待服务启动
    sleep 3

    if systemctl is-active --quiet docker; then
        log_success "Docker 服务重启成功"
    else
        log_error "Docker 服务重启失败"
        exit 1
    fi
}

# 显示可用镜像源
show_available_mirrors() {
    log_title "可用 Docker 镜像源"
    echo
    printf "%-15s %-50s\n" "名称" "地址"
    printf "%-15s %-50s\n" "----" "----------------------------------------"
    for mirror_name in "${MIRROR_NAMES[@]}"; do
        local mirror_var="MIRRORS_${mirror_name}"
        local mirror_url="${!mirror_var}"
        printf "%-15s %-50s\n" "$mirror_name" "$mirror_url"
    done
    echo
}

# 显示当前配置
show_current_config() {
    log_title "当前 Docker 配置"
    echo

    if [[ ! -f "$DOCKER_CONFIG" ]]; then
        echo "暂无配置文件"
        return
    fi

    local config=$(get_current_config)

    # 检查是否有 registry-mirrors
    if echo "$config" | jq -e '.["registry-mirrors"]' >/dev/null 2>&1; then
        local mirrors=$(echo "$config" | jq -r '.["registry-mirrors"][]')
        if [[ -n "$mirrors" ]]; then
            echo "镜像加速器:"
            echo "$mirrors" | while read -r mirror; do
                # 反向查找镜像源名称
                local name="未知"
                for key in "${!MIRRORS[@]}"; do
                    if [[ "${MIRRORS[$key]}" == "$mirror" ]]; then
                        name="$key"
                        break
                    fi
                done
                echo "  - $name ($mirror)"
            done
        else
            echo "镜像加速器: 未配置"
        fi
    else
        echo "镜像加速器: 未配置"
    fi

    # 显示其他配置
    if echo "$config" | jq -e '.log-driver' >/dev/null 2>&1; then
        local log_driver=$(echo "$config" | jq -r '.["log-driver"]')
        echo "日志驱动: $log_driver"
    fi

    if echo "$config" | jq -e '.log-opts' >/dev/null 2>&1; then
        local max_size=$(echo "$config" | jq -r '.["log-opts"]["max-size"] // empty')
        local max_file=$(echo "$config" | jq -r '.["log-opts"]["max-file"] // empty')
        if [[ -n "$max_size" && -n "$max_file" ]]; then
            echo "日志配置: 最大文件大小 $max_size, 保留 $max_file 个文件"
        fi
    fi
    echo
}

# 添加镜像源
add_mirror() {
    local mirror_name="$1"

    local mirror_var="MIRRORS_${mirror_name}"
    if [[ -z "${!mirror_var}" ]]; then
        log_error "未知的镜像源: $mirror_name"
        show_available_mirrors
        exit 1
    fi

    local mirror_url="${!mirror_var}"

    check_config
    backup_config

    local config=$(get_current_config)

    # 检查是否已存在
    if echo "$config" | jq -e ".\"registry-mirrors\"[]? | select(. == \"$mirror_url\")" >/dev/null 2>&1; then
        log_warning "镜像源 $mirror_name 已经配置"
        return
    fi

    # 添加到配置中
    if echo "$config" | jq -e '.["registry-mirrors"]' >/dev/null 2>&1; then
        # 如果已存在 registry-mirrors 数组，添加到数组中
        config=$(echo "$config" | jq --arg url "$mirror_url" '.["registry-mirrors"] += [$url]')
    else
        # 如果不存在，创建新的数组
        config=$(echo "$config" | jq --arg url "$mirror_url" '. + {"registry-mirrors": [$url]}')
    fi

    update_config "$config"
    log_success "已添加镜像源: $mirror_name ($mirror_url)"
}

# 删除镜像源
remove_mirror() {
    local mirror_name="$1"

    if [[ "$mirror_name" == "all" ]]; then
        # 删除所有镜像源
        check_config
        backup_config

        local config=$(get_current_config)
        config=$(echo "$config" | jq 'del(.["registry-mirrors"])')
        update_config "$config"
        log_success "已删除所有镜像源"
        return
    fi

    local mirror_var="MIRRORS_${mirror_name}"
    if [[ -z "${!mirror_var}" ]]; then
        log_error "未知的镜像源: $mirror_name"
        show_available_mirrors
        exit 1
    fi

    local mirror_url="${!mirror_var}"

    check_config
    backup_config

    local config=$(get_current_config)

    # 检查是否存在
    if ! echo "$config" | jq -e ".\"registry-mirrors\"[]? | select(. == \"$mirror_url\")" >/dev/null 2>&1; then
        log_warning "镜像源 $mirror_name 未配置"
        return
    fi

    # 从配置中删除
    config=$(echo "$config" | jq --arg url "$mirror_url" '.["registry-mirrors"] = (.["registry-mirrors"] | map(select(. != $url)))')

    # 如果数组为空，删除整个键
    if echo "$config" | jq -e '.["registry-mirrors"] | length == 0' >/dev/null 2>&1; then
        config=$(echo "$config" | jq 'del(.["registry-mirrors"])')
    fi

    update_config "$config"
    log_success "已删除镜像源: $mirror_name ($mirror_url)"
}

# 设置推荐镜像源组合
set_recommended() {
    log_info "设置推荐的镜像源组合 (阿里云 + 网易 + 华为云)..."

    check_config
    backup_config

    local config=$(get_current_config)

    # 推荐的镜像源组合
    local recommended_mirrors=(
        "https://registry.cn-hangzhou.aliyuncs.com"
        "https://hub-mirror.c.163.com"
        "https://mirror-1.trustie.net"
    )

    # 创建新的配置
    config=$(echo "$config" | jq --argjson mirrors "$(printf '%s\n' "${recommended_mirrors[@]}" | jq -R . | jq -s .)" '. + {"registry-mirrors": $mirrors}')

    update_config "$config"
    log_success "已设置推荐镜像源组合"
}

# 测试镜像源速度
test_mirrors() {
    log_title "测试镜像源速度"
    echo

    local test_image="hello-world:latest"

    # 清理可能存在的测试镜像
    docker rmi "$test_image" >/dev/null 2>&1 || true

    echo "测试镜像: $test_image"
    echo "测试结果:"
    printf "%-15s %-15s %-20s\n" "镜像源" "状态" "耗时"
    printf "%-15s %-15s %-20s\n" "-------" "-----" "----"

    for mirror_name in "${MIRROR_NAMES[@]}"; do
        local mirror_var="MIRRORS_${mirror_name}"
        local mirror_url="${!mirror_var}"

        # 临时设置镜像源
        check_config
        local original_config=$(get_current_config)
        local test_config=$(echo "$original_config" | jq --arg url "$mirror_url" '. + {"registry-mirrors": [$url]}')
        echo "$test_config" > "$DOCKER_CONFIG"

        restart_docker >/dev/null 2>&1

        # 测试拉取镜像
        local start_time=$(date +%s.%3N)
        if timeout 30 docker pull "$test_image" >/dev/null 2>&1; then
            local end_time=$(date +%s.%3N)
            local duration=$(echo "$end_time - $start_time" | bc 2>/dev/null || echo "0")
            printf "%-15s ${GREEN}%-15s${NC} %-20s\n" "$mirror_name" "成功" "${duration}s"
        else
            printf "%-15s ${RED}%-15s${NC} %-20s\n" "$mirror_name" "失败" "超时"
        fi

        # 清理测试镜像
        docker rmi "$test_image" >/dev/null 2>&1 || true
    done

    # 恢复原始配置
    echo "$original_config" > "$DOCKER_CONFIG"
    restart_docker >/dev/null 2>&1

    echo
    log_info "速度测试完成"
}

# 恢复备份配置
restore_config() {
    local backup_file="$1"

    if [[ -z "$backup_file" ]]; then
        # 查找最新的备份文件
        backup_file=$(ls -t /etc/docker/daemon.json.backup.* 2>/dev/null | head -1)
        if [[ -z "$backup_file" ]]; then
            log_error "未找到备份文件"
            exit 1
        fi
    fi

    if [[ ! -f "$backup_file" ]]; then
        log_error "备份文件不存在: $backup_file"
        exit 1
    fi

    cp "$backup_file" "$DOCKER_CONFIG"
    log_success "已恢复配置: $backup_file"
}

# 显示帮助信息
show_help() {
    cat << EOF
Docker 镜像加速器快速设置脚本

此脚本帮助您快速配置 Docker 镜像加速器，提升镜像拉取速度。

用法: $0 [选项] [操作] [参数]

操作:
    list                显示可用镜像源
    show                显示当前配置
    add <镜像源>        添加镜像源 (例: add aliyun)
    remove <镜像源>     删除镜像源 (例: remove aliyun)
    remove all          删除所有镜像源
    recommend           设置推荐镜像源组合
    test                测试各镜像源速度
    restore [文件]      恢复备份配置
    restart             重启 Docker 服务
    demo                演示模式 (无需 root 权限)

选项:
    -h, --help          显示此帮助信息

可用镜像源:
    aliyun     - 阿里云镜像加速器
    tencent    - 腾讯云镜像加速器
    huawei     - 华为云镜像加速器
    netease    - 网易镜像加速器 (推荐)
    ustc       - 中科大镜像加速器
    docker_cn  - Docker 中国官方镜像
    baidu      - 百度云镜像加速器
    tsinghua   - 清华大学镜像加速器 (推荐)

示例:
    $0 demo                          # 演示模式
    sudo $0 list                     # 显示可用镜像源
    sudo $0 show                     # 显示当前配置
    sudo $0 add aliyun               # 添加阿里云镜像源
    sudo $0 add netease              # 添加网易镜像源
    sudo $0 recommend                # 设置推荐组合 (阿里云+网易+华为云)
    sudo $0 test                     # 测试镜像源速度
    sudo $0 remove aliyun            # 删除阿里云镜像源
    sudo $0 remove all               # 删除所有镜像源
    sudo $0 restart                  # 重启 Docker 服务

EOF
}

# 主函数
main() {
    # 检查演示模式
    if [[ "$1" == "demo" ]]; then
        log_title "Docker 镜像加速器设置脚本 - 演示模式"
        echo
        show_available_mirrors
        echo "注意: 实际配置需要 root 权限"
        echo "运行: sudo $0 [操作]"
        echo
        exit 0
    fi

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

    # 检查 jq 是否安装
    if ! command -v jq &> /dev/null; then
        log_info "安装 jq 工具..."
        apt update && apt install -y jq
    fi

    # 检查 bc 是否安装（用于速度测试）
    if ! command -v bc &> /dev/null; then
        apt install -y bc
    fi

    local action="$1"
    shift

    case $action in
        list)
            show_available_mirrors
            ;;
        show)
            check_docker
            show_current_config
            ;;
        add)
            check_docker
            if [[ $# -eq 0 ]]; then
                log_error "请指定要添加的镜像源"
                show_available_mirrors
                exit 1
            fi
            add_mirror "$1"
            restart_docker
            ;;
        remove)
            check_docker
            if [[ $# -eq 0 ]]; then
                log_error "请指定要删除的镜像源"
                exit 1
            fi
            remove_mirror "$1"
            restart_docker
            ;;
        recommend)
            check_docker
            set_recommended
            restart_docker
            ;;
        test)
            check_docker
            test_mirrors
            ;;
        restore)
            restore_config "$1"
            restart_docker
            ;;
        restart)
            check_docker
            restart_docker
            ;;
        ""|-h|--help)
            show_help
            ;;
        *)
            log_error "未知操作: $action"
            show_help
            exit 1
            ;;
    esac

    log_success "操作完成"
}

# 执行主函数
main "$@"