#!/bin/bash
# @desc: DNS 解析延迟测试(剔除国内CDN干扰版)
# @usage: wcl dns
#   不带参数, 直接跑; 依次测国外/国内常见网站

# ================= 域名列表 (已剔除拥有国内CDN的国外大站) =================
# 国外站点：替换了 microsoft/apple，新增 medium, discord, archive.org
FOREIGN_DOMAINS=(
    "google.com"
    "youtube.com"
    "cloudflare.com"
    "github.com"
    "amazon.com"
    "facebook.com"
    "twitter.com"
    "wikipedia.org"
    "medium.com"
    "netflix.com"
    "telegram.org"
    "openai.com"
    "discord.com"
    "twitch.tv"
    "salesforce.com"
    "archive.org"
)

# 国内站点：保持不变
CN_DOMAINS=(
    "baidu.com"
    "qq.com"
    "taobao.com"
    "jd.com"
    "bilibili.com"
    "zhihu.com"
    "aliyun.com"
    "163.com"
    "csdn.net"
    "weibo.com"
    "douyin.com"
    "gitee.com"
    "meituan.com"
    "feishu.cn"
    "tmall.com"
)

# ================= 工具变量 =================
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[0;33m'
NC='\033[0m'

# ================= 核心功能函数 =================

# 1. 检查依赖
function check_dep() {
    if ! command -v dig &> /dev/null; then
        echo -e "${RED}错误: 未找到 'dig' 命令。${NC}"
        echo "请安装 bind-utils (CentOS) 或 dnsutils (Debian/Ubuntu)."
        exit 1
    fi
}

# 2. 暴力清理本地 DNS 缓存
function flush_local_dns() {
    echo -n "正在清理本地 DNS 缓存... "
    local flushed=0

    # (1) systemd-resolved
    if command -v resolvectl &> /dev/null; then
        sudo resolvectl flush-caches 2>/dev/null
        flushed=1
    elif command -v systemd-resolve &> /dev/null; then
        sudo systemd-resolve --flush-caches 2>/dev/null
        flushed=1
    fi

    # (2) nscd
    if systemctl is-active --quiet nscd 2>/dev/null || [ -f /etc/init.d/nscd ]; then
        sudo nscd -i hosts 2>/dev/null || sudo service nscd restart 2>/dev/null
        flushed=1
    fi

    # (3) dnsmasq
    if systemctl is-active --quiet dnsmasq 2>/dev/null; then
        sudo systemctl restart dnsmasq 2>/dev/null
        flushed=1
    fi
    
    # (4) NetworkManager
    if systemctl is-active --quiet NetworkManager 2>/dev/null; then
        sudo systemctl reload NetworkManager 2>/dev/null
        flushed=1
    fi

    if [ $flushed -eq 1 ]; then
        echo -e "${GREEN}完成${NC}"
    else
        echo -e "${YELLOW}未检测到常见缓存服务，跳过${NC}"
    fi
}

# 3. 单个域名测试
function test_latency() {
    local domain=$1
    # dig 参数：+stats 只显示统计，+tries=1 不重试，+time=2 超时2秒
    local t=$(dig "$domain" +noall +stats +tries=1 +time=2 | grep "Query time" | awk '{print $4}')
    
    if [[ -n "$t" && "$t" =~ ^[0-9]+$ ]]; then
        echo "$t"
    else
        echo "FAIL"
    fi
}

# 4. 批量测试逻辑
function run_benchmark() {
    local list_name=$1
    shift
    local domains=("$@")
    
    local total_time=0
    local success_count=0
    
    echo -e "\n### 测试组: $list_name"
    printf "%-25s | %-10s\n" "Domain" "Latency"
    echo "----------------------------------------"
    
    # 每组测试前清理缓存
    flush_local_dns
    
    for domain in "${domains[@]}"; do
        local result=$(test_latency "$domain")
        
        if [ "$result" == "FAIL" ]; then
            printf "%-25s | ${RED}%-10s${NC}\n" "$domain" "Timeout"
        else
            local color=$NC
            if [ "$result" -lt 20 ]; then color=$GREEN; fi     # 极快
            if [ "$result" -gt 200 ]; then color=$RED; fi      # 较慢
            
            printf "%-25s | ${color}%-8s ms${NC}\n" "$domain" "$result"
            
            total_time=$((total_time + result))
            success_count=$((success_count + 1))
        fi
        sleep 0.2
    done
    
    echo "----------------------------------------"
    if [ "$success_count" -gt 0 ]; then
        local avg=$((total_time / success_count))
        echo -e ">>> $list_name 平均延迟: ${YELLOW}${avg} ms${NC}"
    else
        echo -e ">>> $list_name 全部超时/失败"
    fi
}

# ================= 主程序 =================
check_dep

if [ "$EUID" -ne 0 ]; then
    echo "提示: 清理本地 DNS 缓存可能需要 root 权限..."
    # sudo -v # 这一行注释掉，避免脚本卡在后台等待，让具体的sudo命令去触发提示
fi

clear
echo "========================================"
echo "   DNS 解析延迟测试 (剔除国内CDN干扰版)   "
echo "========================================"

run_benchmark "国外常见网站 (Foreign)" "${FOREIGN_DOMAINS[@]}"
run_benchmark "国内常见网站 (Domestic)" "${CN_DOMAINS[@]}"

echo -e "\n测试全部结束。"