#!/bin/sh
#
# api-speed-test.sh — 测试并对比两个 TokenFlows API 域名的连通性与连接速度
#
# 用法:
#   ./api-speed-test.sh [选项]
#   curl -fsSL https://tokenflows.org/api-speed-test.sh | sh -s -- [选项]
#
# 选项:
#   -n, --rounds N      每个域名测试轮数 (默认 5)
#   -p, --path PATH     请求路径 (默认 /)
#   -t, --timeout SEC   单次请求超时秒数 (默认 10)
#       --no-ping       跳过 ping 测试 (部分网络环境禁用 ICMP)
#   -h, --help          显示帮助
#
# 退出码:
#   0  两个域名均连通并完成对比
#   1  存在无法连通的域名
#   2  参数错误或缺少依赖
#
set -u

DOMAIN_NEW="https://api-new.tokenflows.org"
DOMAIN_US="https://api-us.tokenflows.org"

ROUNDS=5
TIMEOUT=10
REQ_PATH="/"
DO_PING=1
PING_COUNT=5
UA="tokenflows-api-check/1.0"

usage() {
    cat <<'USAGE'
api-speed-test.sh — 测试并对比两个 TokenFlows API 域名的连通性与连接速度

用法:
  ./api-speed-test.sh [选项]
  curl -fsSL https://tokenflows.org/api-speed-test.sh | sh -s -- [选项]

选项:
  -n, --rounds N      每个域名测试轮数 (默认 5)
  -p, --path PATH     请求路径 (默认 /)
  -t, --timeout SEC   单次请求超时秒数 (默认 10)
      --no-ping       跳过 ping 测试 (部分网络环境禁用 ICMP)
  -h, --help          显示帮助

退出码:
  0  两个域名均连通并完成对比
  1  存在无法连通的域名
  2  参数错误或缺少依赖
USAGE
}

# ---------- 参数解析 ----------
while [ $# -gt 0 ]; do
    case "$1" in
        -n|--rounds)  ROUNDS="${2:-}"; shift 2 ;;
        -p|--path)    REQ_PATH="${2:-}"; shift 2 ;;
        -t|--timeout) TIMEOUT="${2:-}"; shift 2 ;;
        --no-ping)    DO_PING=0; shift ;;
        -h|--help)    usage; exit 0 ;;
        *) echo "未知参数: $1" >&2; usage >&2; exit 2 ;;
    esac
done

case "$ROUNDS" in
    ''|*[!0-9]*) echo "错误: --rounds 必须是正整数 (当前: $ROUNDS)" >&2; exit 2 ;;
esac
case "$TIMEOUT" in
    ''|*[!0-9]*) echo "错误: --timeout 必须是正整数 (当前: $TIMEOUT)" >&2; exit 2 ;;
esac
[ "$ROUNDS" -ge 1 ] || { echo "错误: --rounds 至少为 1" >&2; exit 2; }
case "$REQ_PATH" in
    /*) ;;
    *) REQ_PATH="/$REQ_PATH" ;;
esac

command -v curl >/dev/null 2>&1 || { echo "错误: 缺少 curl" >&2; exit 2; }
command -v awk  >/dev/null 2>&1 || { echo "错误: 缺少 awk" >&2; exit 2; }

# ---------- 颜色 ----------
if [ -t 1 ]; then
    C_RESET=$(printf '\033[0m')
    C_BOLD=$(printf '\033[1m')
    C_DIM=$(printf '\033[2m')
    C_RED=$(printf '\033[31m')
    C_GREEN=$(printf '\033[32m')
    C_YELLOW=$(printf '\033[33m')
    C_CYAN=$(printf '\033[36m')
else
    C_RESET=""; C_BOLD=""; C_DIM=""; C_RED=""; C_GREEN=""; C_YELLOW=""; C_CYAN=""
fi

TMP=$(mktemp -d "${TMPDIR:-/tmp}/api-speed-test.XXXXXX") || exit 2
trap 'rm -rf "$TMP"' EXIT INT TERM

DOMAINS="$DOMAIN_NEW $DOMAIN_US"
IDX=0
for d in $DOMAINS; do
    IDX=$((IDX + 1))
    echo "$d" > "$TMP/url.$IDX"
done
TOTAL_DOMAINS=$IDX

# ---------- 工具函数 ----------
host_of() { echo "$1" | sed -e 's|^https\{0,1\}://||' -e 's|/.*$||'; }

fmt_ms() {
    # 输入: 秒; 输出: 毫秒(整数)
    awk -v v="${1:-0}" 'BEGIN { printf "%d", (v + 0) * 1000 }'
}

kv() {
    # kv <文件> <键>
    [ -f "$1" ] || return 0
    sed -n "s/^$2=//p" "$1" | head -n 1
}

color_by_threshold() {
    # color_by_threshold <毫秒值> <good> <warn>  -> 输出着色后的字符串
    _v=$1
    if [ "$_v" -le "$2" ] 2>/dev/null; then
        printf '%s' "${C_GREEN}${_v}ms${C_RESET}"
    elif [ "$_v" -le "$3" ] 2>/dev/null; then
        printf '%s' "${C_YELLOW}${_v}ms${C_RESET}"
    else
        printf '%s' "${C_RED}${_v}ms${C_RESET}"
    fi
}

# ---------- 开始 ----------
printf '%s\n' "${C_BOLD}============================================================${C_RESET}"
printf '%s\n' "${C_BOLD} TokenFlows API 域名连通性与速度测试${C_RESET}"
printf '%s\n' "${C_BOLD}============================================================${C_RESET}"
printf ' 请求路径 : %s\n' "$REQ_PATH"
printf ' 测试轮数 : %s 轮/域名\n' "$ROUNDS"
printf ' 单次超时 : %ss\n' "$TIMEOUT"
printf ' 测试时间 : %s\n' "$(date '+%Y-%m-%d %H:%M:%S')"
printf ' 出口 IP  : '

if command -v curl >/dev/null 2>&1; then
    _ip=$(curl -s --max-time "$TIMEOUT" https://api.ipify.org 2>/dev/null || true)
    if [ -n "$_ip" ]; then printf '%s' "$_ip"; else printf '未知'; fi
    unset _ip
fi
printf '\n\n'

# ---------- 1. DNS 解析 ----------
printf '%s\n' "${C_BOLD}[1/4] DNS 解析${C_RESET}"
for i in 1 2; do
    url=$(cat "$TMP/url.$i")
    host=$(host_of "$url")
    if command -v dig >/dev/null 2>&1; then
        ips=$(dig +short +time=3 +tries=1 A "$host" 2>/dev/null | grep -E '^[0-9]+\.' | tr '\n' ' ')
    elif command -v nslookup >/dev/null 2>&1; then
        ips=$(nslookup "$host" 2>/dev/null | awk '/^Address: /{print $2}' | tr '\n' ' ')
    else
        ips=""
    fi
    if [ -n "$(echo "$ips" | tr -d ' ')" ]; then
        printf '  %-34s %s%s%s\n' "$host" "$C_DIM" "$ips" "$C_RESET"
    else
        printf '  %-34s %s(解析失败或无法查询 DNS)%s\n' "$host" "$C_RED" "$C_RESET"
    fi
done
printf '\n'

# ---------- 2. ICMP ping ----------
if [ "$DO_PING" -eq 1 ] && command -v ping >/dev/null 2>&1; then
    printf '%s\n' "${C_BOLD}[2/4] ICMP Ping (延迟参考)${C_RESET}"
    for i in 1 2; do
        url=$(cat "$TMP/url.$i")
        host=$(host_of "$url")
        out=$(ping -c "$PING_COUNT" -q "$host" 2>/dev/null || true)
        stats=$(echo "$out" | grep -E 'min/avg/max' | sed 's|.*= *||')
        if [ -n "$stats" ]; then
            avg=$(echo "$stats" | awk -F'/' '{print $2}')
            loss=$(echo "$out" | sed -n 's/.*, \([0-9.]*\)% packet loss.*/\1/p')
            [ -n "$loss" ] || loss="?"
            printf '  %-34s avg %-9s 丢包 %s%%\n' "$host" "${avg}ms" "$loss"
        else
            printf '  %-34s %s无响应 (ICMP 被屏蔽或网络不通)%s\n' "$host" "$C_YELLOW" "$C_RESET"
        fi
    done
    printf '\n'
else
    printf '%s\n\n' "${C_DIM}[2/4] ICMP Ping 已跳过${C_RESET}"
fi

# ---------- 3. HTTPS 连接测试 ----------
printf '%s\n' "${C_BOLD}[3/4] HTTPS 连接测试 (DNS / TCP / TLS / TTFB)${C_RESET}"
for i in 1 2; do
    url=$(cat "$TMP/url.$i")
    host=$(host_of "$url")
    tsv="$TMP/run.$i.tsv"
    : > "$tsv"
    printf '  %s%s%s\n' "$C_CYAN" "$url$REQ_PATH" "$C_RESET"

    r=1
    while [ "$r" -le "$ROUNDS" ]; do
        raw=$(curl -sS -o /dev/null \
                   -A "$UA" \
                   --max-time "$TIMEOUT" \
                   -w '%{http_code} %{time_namelookup} %{time_connect} %{time_appconnect} %{time_starttransfer} %{time_total} %{size_download}' \
                   "$url$REQ_PATH" 2>"$TMP/err.$i.$r")
        rc=$?

        if [ "$rc" -ne 0 ]; then
            err=$(tr '\n' ' ' < "$TMP/err.$i.$r" | cut -c1-70)
            [ -n "$err" ] || err="curl 退出码 $rc"
            printf '    %s第 %s/%s 次  失败: %s%s\n' "$C_RED" "$r" "$ROUNDS" "$err" "$C_RESET"
            printf '000\t0\t0\t0\t0\t0\t0\t%s\t%s\n' "$rc" "$err" >> "$tsv"
        else
            code=$(echo "$raw" | awk '{print $1}')
            nl=$(echo "$raw" | awk '{print $2}')
            cn=$(echo "$raw" | awk '{print $3}')
            ac=$(echo "$raw" | awk '{print $4}')
            st=$(echo "$raw" | awk '{print $5}')
            tt=$(echo "$raw" | awk '{print $6}')
            sz=$(echo "$raw" | awk '{print $7}')
            printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t0\t\n' \
                "$code" "$nl" "$cn" "$ac" "$st" "$tt" "$sz" >> "$tsv"

            d_dns=$(fmt_ms "$nl")
            d_tcp=$(fmt_ms "$(awk -v a="$cn" -v b="$nl" 'BEGIN{d=a-b; if(d<0)d=0; print d}')")
            d_tls=$(fmt_ms "$(awk -v a="$ac" -v b="$cn" 'BEGIN{d=a-b; if(d<0)d=0; print d}')")
            d_ttfb=$(fmt_ms "$(awk -v a="$st" -v b="$ac" 'BEGIN{d=a-b; if(d<0)d=0; print d}')")
            d_total=$(fmt_ms "$tt")
            if [ "$sz" -ge 1024 ] 2>/dev/null; then
                sz_txt=$(awk -v s="$sz" 'BEGIN{printf "%.1fKB", s/1024}')
            else
                sz_txt="${sz}B"
            fi

            case "$code" in
                2*|3*) code_txt="${C_GREEN}HTTP $code${C_RESET}" ;;
                4*)    code_txt="${C_YELLOW}HTTP $code${C_RESET}" ;;
                *)     code_txt="${C_RED}HTTP $code${C_RESET}" ;;
            esac

            printf '    %s第 %s/%s 次%s  %s  DNS %s  TCP %s  TLS %s  TTFB %s  总计 %s%s%s\n' \
                "$C_DIM" "$r" "$ROUNDS" "$C_RESET" "$code_txt" \
                "$(color_by_threshold "$d_dns" 200 800)" \
                "$(color_by_threshold "$d_tcp" 200 800)" \
                "$(color_by_threshold "$d_tls" 400 1500)" \
                "$(color_by_threshold "$d_ttfb" 500 2000)" \
                "$C_BOLD" "$(color_by_threshold "$d_total" 1000 3000)" "$C_RESET"
        fi
        r=$((r + 1))
    done
    printf '\n'
done

# ---------- 4. 汇总与对比 ----------
summarize() {
    awk -F'\t' '
    function median(arr, n,   i, j, tmp) {
        if (n == 0) return 0
        for (i = 2; i <= n; i++) {
            tmp = arr[i]; j = i - 1
            while (j >= 1 && arr[j] > tmp) { arr[j+1] = arr[j]; j-- }
            arr[j+1] = tmp
        }
        if (n % 2 == 1) return arr[int(n/2)+1]
        return (arr[n/2] + arr[n/2+1]) / 2
    }
    {
        code = $1; nl = $2; cn = $3; ac = $4; st = $5; tt = $6; sz = $7; rc = $8
        rounds++
        if (rc + 0 == 0 && code + 0 > 0) {
            n++
            dns = nl * 1000
            tcp = (cn - nl) * 1000; if (tcp < 0) tcp = 0
            tls = (ac - cn) * 1000; if (tls < 0) tls = 0
            ttfb = (st - ac) * 1000; if (ttfb < 0) ttfb = 0
            total = tt * 1000
            v_dns[n] = dns; v_tcp[n] = tcp; v_tls[n] = tls; v_ttfb[n] = ttfb; v_total[n] = total
            s_dns += dns; s_tcp += tcp; s_tls += tls; s_ttfb += ttfb; s_total += total; s_size += sz
            if (n == 1 || total < min) min = total
            if (n == 1 || total > max) max = total
            codes[code]++
        } else {
            fails++
        }
    }
    END {
        printf "ROUNDS=%d\n", rounds
        printf "OK=%d\n", n
        printf "FAIL=%d\n", fails
        c = ""
        for (k in codes) { if (c != "") c = c " "; c = c k "x" codes[k] }
        printf "CODES=%s\n", (c == "" ? "-" : c)
        if (n == 0) { printf "REACHABLE=0\n"; exit }
        printf "REACHABLE=1\n"
        printf "DNS_AVG=%.1f\n", s_dns / n
        printf "TCP_AVG=%.1f\n", s_tcp / n
        printf "TLS_AVG=%.1f\n", s_tls / n
        printf "TTFB_AVG=%.1f\n", s_ttfb / n
        printf "TOTAL_AVG=%.1f\n", s_total / n
        printf "TOTAL_MIN=%.1f\n", min
        printf "TOTAL_MAX=%.1f\n", max
        printf "TOTAL_MED=%.1f\n", median(v_total, n)
        printf "TTFB_MED=%.1f\n", median(v_ttfb, n)
        printf "SIZE_AVG=%.0f\n", s_size / n
    }
    ' "$1" > "$2"
}

printf '%s\n' "${C_BOLD}============================================================${C_RESET}"
printf '%s\n' "${C_BOLD} 汇总结果${C_RESET}"
printf '%s\n' "${C_BOLD}============================================================${C_RESET}"
printf ' %-32s %-6s %-10s %-9s %-9s %-9s %-9s %-9s %-9s %-9s\n' \
    "域名" "成功" "HTTP" "DNS" "TCP" "TLS" "TTFB" "总计" "最快" "最慢"

for i in 1 2; do
    url=$(cat "$TMP/url.$i")
    host=$(host_of "$url")
    summarize "$TMP/run.$i.tsv" "$TMP/sum.$i"

    reachable=$(kv "$TMP/sum.$i" REACHABLE)
    ok=$(kv "$TMP/sum.$i" OK)
    rounds=$(kv "$TMP/sum.$i" ROUNDS)
    codes=$(kv "$TMP/sum.$i" CODES)

    if [ "$reachable" = "1" ]; then
        printf ' %-32s %-6s %-10s %-9s %-9s %-9s %-9s %-9s %-9s %-9s\n' \
            "$host" \
            "$ok/$rounds" \
            "$codes" \
            "$(kv "$TMP/sum.$i" DNS_AVG)ms" \
            "$(kv "$TMP/sum.$i" TCP_AVG)ms" \
            "$(kv "$TMP/sum.$i" TLS_AVG)ms" \
            "$(kv "$TMP/sum.$i" TTFB_AVG)ms" \
            "$(kv "$TMP/sum.$i" TOTAL_AVG)ms" \
            "$(kv "$TMP/sum.$i" TOTAL_MIN)ms" \
            "$(kv "$TMP/sum.$i" TOTAL_MAX)ms"
    else
        printf ' %-32s %-6s %-10s %s\n' "$host" "$ok/$rounds" "$codes" "${C_RED}无法连通${C_RESET}"
    fi
done
printf '\n'

# ---------- 结论 ----------
R1=$(kv "$TMP/sum.1" REACHABLE)
R2=$(kv "$TMP/sum.2" REACHABLE)
H1=$(host_of "$(cat "$TMP/url.1")")
H2=$(host_of "$(cat "$TMP/url.2")")

printf '%s\n' "${C_BOLD}结论${C_RESET}"
if [ "$R1" = "1" ] && [ "$R2" = "1" ]; then
    T1=$(kv "$TMP/sum.1" TOTAL_AVG)
    T2=$(kv "$TMP/sum.2" TOTAL_AVG)
    F1=$(kv "$TMP/sum.1" TTFB_MED)
    F2=$(kv "$TMP/sum.2" TTFB_MED)
    F1x=$(kv "$TMP/sum.1" FAIL)
    F2x=$(kv "$TMP/sum.2" FAIL)

    awk -v h1="$H1" -v h2="$H2" -v t1="$T1" -v t2="$T2" -v f1="$F1" -v f2="$F2" -v x1="$F1x" -v x2="$F2x" -v g="$C_GREEN" -v y="$C_YELLOW" -v r="$C_RED" -v z="$C_RESET" 'BEGIN {
        winner = ""
        if (t1 < t2) winner = h1; else if (t2 < t1) winner = h2
        diff = 0
        if (t1 > 0 && t2 > 0) {
            lo = (t1 < t2 ? t1 : t2); hi = (t1 > t2 ? t1 : t2)
            diff = (hi - lo) / hi * 100
        }
        printf "  平均总耗时 : %s %.0fms  vs  %s %.0fms\n", h1, t1, h2, t2
        printf "  TTFB 中位数: %s %.0fms  vs  %s %.0fms\n", h1, f1, h2, f2
        if (x1 + 0 > 0 || x2 + 0 > 0)
            printf "  失败轮次   : %s %d 次, %s %d 次\n", h1, x1, h2, x2
        if (winner == "") {
            printf "  两者速度基本相当 (差异 < 1%%)\n"
        } else if (diff < 10) {
            printf "  %s差异小于 10%%，可视为网络抖动，两者速度相当%s\n", y, z
        } else {
            printf "  %s推荐: %s%s  (快 %.1f%%)\n", g, winner, z, diff
        }
    }'
    if [ "$F1x" != "0" ] || [ "$F2x" != "0" ]; then
        printf '  %s注意: 存在请求失败，建议结合失败域名排查%s\n' "$C_YELLOW" "$C_RESET"
    fi
    printf '  建议: 客户端优先使用低延迟域名，并保留另一域名作为自动切换的备份。\n'
    exit 0
elif [ "$R1" = "1" ]; then
    printf '  %s%s%s 可正常连通；%s%s%s 无法连通，请检查网络/DNS/服务状态。\n' \
        "$C_GREEN" "$H1" "$C_RESET" "$C_RED" "$H2" "$C_RESET"
    printf '  %s推荐: %s%s\n' "$C_GREEN" "$H1" "$C_RESET"
    exit 1
elif [ "$R2" = "1" ]; then
    printf '  %s%s%s 可正常连通；%s%s%s 无法连通，请检查网络/DNS/服务状态。\n' \
        "$C_GREEN" "$H2" "$C_RESET" "$C_RED" "$H1" "$C_RESET"
    printf '  %s推荐: %s%s\n' "$C_GREEN" "$H2" "$C_RESET"
    exit 1
else
    printf '  %s两个域名均无法连通，请检查本机网络、DNS 或代理设置。%s\n' "$C_RED" "$C_RESET"
    exit 1
fi