#!/bin/bash

# ==============================================================================
#  test-npm-registry.sh
#
#  一个用于测试不同 npm registry 速度和延迟的 Shell 脚本。
#  它会测试延迟、清理缓存、安装一个包来衡量速度，并最后提供一份总结报告。
# ==============================================================================

# --- 配置 ---

# 你想用来测试的包 (选择一个体积适中、无复杂依赖的包)
# lodash 是个不错的选择，体积小，无依赖。
TEST_PACKAGE="lodash"

# 要测试的 registry 列表 (可以随意添加或删除)
REGISTRIES=(
    "https://registry.npmjs.org/" # 官方
    "http://registry.npmmirror.com" # 国内镜像
    "https://npm.aliyun.com" # 阿里云
    "https://mirrors.cloud.tencent.com/npm/" # 腾讯云
    "https://mirrors.huaweicloud.com/repository/npm/" # 华为云
    "https://mirrors.163.com/npm/" # 网易
    "http://mirrors.ustc.edu.cn/" # 中科大
    "https://mirrors.tuna.tsinghua.edu.cn/" # 清华
)

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

# --- 检查依赖 ---
command -v npm >/dev/null 2>&1 || { echo -e "${RED}错误: 'npm' 命令未找到。请先安装 Node.js 和 npm。${NC}"; exit 1; }
command -v curl >/dev/null 2>&1 || { echo -e "${RED}错误: 'curl' 命令未找到。请先安装 curl。${NC}"; exit 1; }


# --- 主要功能 ---

# 捕获脚本退出信号，确保执行清理函数
trap 'cleanup' EXIT

# 清理函数：恢复原始的 registry 配置
cleanup() {
    if [ -n "$ORIGINAL_REGISTRY" ]; then
        echo -e "\n${YELLOW}恢复原始 npm registry: ${BLUE}$ORIGINAL_REGISTRY${NC}"
        npm config set registry "$ORIGINAL_REGISTRY" > /dev/null
    fi
    # 清理测试产生的临时文件
    rm -f package.json package-lock.json
    rm -rf node_modules/
    echo -e "${GREEN}清理完成。${NC}"
}

# 保存用户当前的 registry
ORIGINAL_REGISTRY=$(npm config get registry)
echo -e "${BLUE}你当前的 npm registry 是: ${YELLOW}$ORIGINAL_REGISTRY${NC}"
echo -e "${BLUE}测试将在结束后自动恢复此设置。${NC}"

# 存储结果的数组
results=()

# 主循环，遍历所有待测试的 registry
for registry in "${REGISTRIES[@]}"; do
    echo -e "\n${BLUE}==============================================================${NC}"
    echo -e "${YELLOW}正在测试源: ${GREEN}${registry}${NC}"
    echo -e "${BLUE}==============================================================${NC}"

    # 1. 测试网络延迟 (Latency)
    # 使用 curl 测试连接、TLS 握手和服务器响应时间
    echo -e "${YELLOW}1. 正在测试网络延迟...${NC}"
    # curl -w: 定义输出格式; -o /dev/null: 不输出下载内容; -s: 静默模式; --max-time: 设置30秒超时
    latency_info=$(curl -w "time_connect=%{time_connect}s | time_starttransfer=%{time_starttransfer}s | time_total=%{time_total}s" -o /dev/null -s --max-time 30 "${registry}${TEST_PACKAGE}")
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}延迟测试成功: ${latency_info}${NC}"
        latency=$(echo $latency_info | awk -F'|' '{print $2}' | sed 's/time_starttransfer=//;s/s//')
    else
        echo -e "${RED}延迟测试失败，无法连接到 ${registry} (30秒超时)${NC}"
        latency="N/A"
    fi

    # 2. 测试安装速度 (Install Speed)
    echo -e "\n${YELLOW}2. 正在测试包安装速度 (${TEST_PACKAGE})...${NC}"

    # 切换到当前要测试的 registry
    npm config set registry "$registry"

    # 清理环境，确保测试公平
    echo "   - 清理缓存和旧文件..."
    rm -rf node_modules/ package-lock.json yarn.lock > /dev/null 2>&1
    npm cache clean --force > /dev/null 2>&1

    # 记录开始时间
    start_time=$(date +%s.%N)

    # 执行 npm install，设置30秒超时
    npm config set fetch-retry-maxtimeout 30000
    install_output=$(npm install "${TEST_PACKAGE}" 2>&1)

    if [ $? -eq 0 ]; then
        # 记录结束时间
        end_time=$(date +%s.%N)
        # 计算耗时
        duration=$(echo "$end_time - $start_time" | bc)
        echo -e "${GREEN}安装成功！耗时: ${duration} 秒${NC}"
        install_time=$duration
    else
        # 检查是否是因为超时导致的失败
        if echo "$install_output" | grep -q "timeout"; then
            echo -e "${RED}安装超时! (30秒超时)${NC}"
            install_time="Timeout"
        else
            echo -e "${RED}安装失败!${NC}"
            echo -e "${RED}错误信息: ${install_output}${NC}"
            install_time="Failed"
        fi
    fi

    # 将结果存入数组，用分号分隔
    results+=("${registry};${latency};${install_time}")

    # 每次测试后都清理一次，避免影响下一次
    rm -rf node_modules/ package-lock.json
done

# --- 结果总结 ---

echo -e "\n\n${BLUE}========================= 测试结果总结 ========================${NC}"
# 打印表头
printf "%-60s %-20s %-20s\n" "Registry 源" "延迟 (s)" "安装耗时 (s)"
printf "%-60s %-20s %-20s\n" "------------------------------------------------------------" "--------------------" "--------------------"

# 对结果进行排序（按安装时间升序），并打印
# sort -t';' -k3 -n 表示以分号为分隔符，按第3个字段（安装时间）进行数值排序
(
for result in "${results[@]}"; do
    # 使用 IFS (Internal Field Separator) 来分割字符串
    IFS=';' read -r reg lat inst <<< "$result"
    # 格式化输出
    printf "%-60s %-20s %-20s\n" "$reg" "$lat" "$inst"
done
) | sort -t';' -k3 -n | while IFS=';' read -r reg lat inst; do
    # 重新读取排序后的行并格式化打印
    printf "%-60s %-20s %-20s\n" "$reg" "$lat" "$inst"
done

echo -e "${BLUE}==============================================================${NC}"
echo -e "${GREEN}测试完成！安装耗时越短越好。${NC}"

# cleanup 函数会在脚本退出时自动被调用，无需手动调用
exit 0

执行以下命令，即可测试 pip 镜像源的下行速度：

bash test-npm-registry.sh