#!/bin/bash

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

echo -e "${BLUE}🐳 Docker镜像源自动测速脚本${NC}"
echo "=================================="

# 定义测试的镜像源列表
declare -A MIRRORS=(
    ["阿里云"]="mirrors.aliyun.com"
    ["中科大"]="mirrors.ustc.edu.cn" 
    ["清华"]="mirrors.tuna.tsinghua.edu.cn"
    ["华为云"]="repo.huaweicloud.com"
    ["腾讯云"]="mirrors.cloud.tencent.com"
    ["网易"]="mirrors.163.com"
)

# 测试文件（使用一个小文件进行测速）
TEST_FILE="dists/bookworm/InRelease"
# 或者使用Docker的实际包文件测试
# TEST_FILE="linux/debian/dists/bookworm/InRelease"

# 存储测速结果
declare -A SPEED_RESULTS
declare -A TIME_RESULTS

# 测速函数
speed_test() {
    local mirror_name=$1
    local mirror_url=$2
    
    echo -e "${YELLOW}测试 ${mirror_name} (${mirror_url})...${NC}"
    
    # 使用curl进行测速，下载一个小文件
    start_time=$(date +%s.%N)
    
    # 测试连接时间和下载速度
    result=$(curl -L --connect-timeout 5 --max-time 10 -o /dev/null -s -w \
    "time_connect:%{time_connect}\ntime_starttransfer:%{time_starttransfer}\ntime_total:%{time_total}\nspeed_download:%{speed_download}" \
    "https://${mirror_url}/${TEST_FILE}" 2>/dev/null)
    
    end_time=$(date +%s.%N)
    
    # 提取速度信息（字节/秒）
    speed_download=$(echo "$result" | grep "speed_download" | cut -d':' -f2)
    time_total=$(echo "$result" | grep "time_total" | cut -d':' -f2)
    
    # 转换为KB/s
    if [[ $speed_download =~ ^[0-9]+([.][0-9]+)?$ ]]; then
        speed_kbs=$(echo "scale=2; $speed_download / 1024" | bc)
    else
        speed_kbs=0
    fi
    
    # 计算响应时间（毫秒）
    response_time=$(echo "scale=2; $time_total * 1000" | bc 2>/dev/null || echo "0")
    
    # 存储结果
    SPEED_RESULTS["$mirror_name"]=$speed_kbs
    TIME_RESULTS["$mirror_name"]=$response_time
    
    if [ $(echo "$speed_kbs > 0" | bc) -eq 1 ]; then
        echo -e "${GREEN}  ✓ 速度: ${speed_kbs} KB/s, 响应时间: ${response_time}ms${NC}"
    else
        echo -e "${RED}  ✗ 连接失败或超时${NC}"
        SPEED_RESULTS["$mirror_name"]=0
        TIME_RESULTS["$mirror_name"]=9999
    fi
    
    echo ""
}

# 更换源的函数
change_mirror() {
    local fastest_mirror=$1
    local fastest_url=$2
    
    echo -e "${GREEN}🎯 切换到最快的镜像源: ${fastest_mirror}${NC}"
    
    # 备份原配置
    if [ -f "/etc/apt/sources.list.d/docker.list" ]; then
        sudo cp /etc/apt/sources.list.d/docker.list /etc/apt/sources.list.d/docker.list.bak.$(date +%Y%m%d_%H%M%S)
    fi
    
    # 创建新的Docker源配置
    sudo tee /etc/apt/sources.list.d/docker.list > /dev/null <<EOF
deb [arch=amd64] https://${fastest_url}/docker-ce/linux/debian bookworm stable
EOF
    
    echo -e "${GREEN}✅ 已更新Docker镜像源为: ${fastest_mirror}${NC}"
    echo -e "${BLUE}🔄 更新软件包列表...${NC}"
    
    sudo apt update
}

# 主函数
main() {
    echo -e "${BLUE}📡 开始测试各镜像源速度...${NC}"
    echo ""
    
    # 测试所有镜像源
    for mirror_name in "${!MIRRORS[@]}"; do
        speed_test "$mirror_name" "${MIRRORS[$mirror_name]}"
        sleep 1  # 避免请求过于频繁
    done
    
    # 找出最快的镜像源
    fastest_mirror=""
    fastest_speed=0
    fastest_time=9999
    
    echo -e "${BLUE}📊 测速结果汇总:${NC}"
    echo "=================================="
    
    for mirror_name in "${!SPEED_RESULTS[@]}"; do
        speed=${SPEED_RESULTS[$mirror_name]}
        time=${TIME_RESULTS[$mirror_name]}
        
        if [ $(echo "$speed > 0" | bc) -eq 1 ]; then
            echo -e "  ${mirror_name}: ${speed} KB/s, ${time}ms"
            
            # 选择策略：优先速度，速度相近时选择响应时间短的
            if [ $(echo "$speed > $fastest_speed" | bc) -eq 1 ] || \
               [ $(echo "$speed == $fastest_speed && $time < $fastest_time" | bc) -eq 1 ]; then
                fastest_speed=$speed
                fastest_time=$time
                fastest_mirror=$mirror_name
            fi
        fi
    done
    
    echo ""
    
    if [ -n "$fastest_mirror" ] && [ $(echo "$fastest_speed > 0" | bc) -eq 1 ]; then
        echo -e "${GREEN}🚀 最快的镜像源: ${fastest_mirror}${NC}"
        echo -e "${GREEN}📈 速度: ${fastest_speed} KB/s, 响应时间: ${fastest_time}ms${NC}"
        echo ""
        
        read -p "是否立即更换为此镜像源？(y/N): " confirm
        if [[ $confirm =~ ^[Yy]$ ]]; then
            change_mirror "$fastest_mirror" "${MIRRORS[$fastest_mirror]}"
        else
            echo -e "${YELLOW}未更换镜像源。${NC}"
        fi
    else
        echo -e "${RED}❌ 所有镜像源测试失败，请检查网络连接。${NC}"
    fi
}

# 检查依赖
check_dependencies() {
    if ! command -v curl &> /dev/null; then
        echo -e "${RED}❌ 需要安装 curl${NC}"
        sudo apt update && sudo apt install -y curl
    fi
    
    if ! command -v bc &> /dev/null; then
        echo -e "${RED}❌ 需要安装 bc${NC}"
        sudo apt update && sudo apt install -y bc
    fi
}

# 运行脚本
check_dependencies
main