#!/bin/bash

# Combined Network Benchmark Script (Resume/Retry Mode)
# Uses jq to skip already successful tests from network_benchmark.json

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
PLAIN='\033[0m'

# Files
JSON_FILE="network_benchmark.json"
TEMP_DIR=$(mktemp -d)
trap 'rm -rf "$TEMP_DIR"' EXIT

# --- 1. Dependencies & Setup ---

echo -e "${YELLOW}Initializing Benchmark...${PLAIN}"

# Check tools
check_cmd() {
    if ! command -v "$1" &> /dev/null; then
        echo -e "${RED}Error: $1 is required but not installed. Please install it (e.g., apt install $1).${PLAIN}"
        exit 1
    fi
}
check_cmd curl
check_cmd grep
check_cmd awk
check_cmd jq  # Added jq requirement

# Load existing JSON if available
PREV_JSON="{}"
if [[ -f "$JSON_FILE" ]]; then
    echo -e "${CYAN}Found existing $JSON_FILE, loading previous results...${PLAIN}"
    # Validate JSON content
    if jq -e . "$JSON_FILE" >/dev/null 2>&1; then
        PREV_JSON=$(cat "$JSON_FILE")
    else
        echo -e "${RED}Existing JSON is invalid, starting fresh.${PLAIN}"
    fi
fi

# Detect Host IP
echo -e "${YELLOW}Detecting Host IP addresses...${PLAIN}"
HOST_IPV4=$(curl -s -4 --connect-timeout 5 icanhazip.com || echo "N/A")
HOST_IPV6=$(curl -s -6 --connect-timeout 5 icanhazip.com || echo "N/A")
echo -e "IPv4: ${GREEN}$HOST_IPV4${PLAIN}"
echo -e "IPv6: ${GREEN}$HOST_IPV6${PLAIN}"

# Determine Architecture for Binaries
ARCH=$(uname -m)
case $ARCH in
    x86_64)  
        YABS_ARCH="x64"
        NWS_ARCH="x86_64" 
        ;;
    i*86)    
        YABS_ARCH="x86"
        NWS_ARCH="i386"
        ;;
    aarch64|armv8*) 
        YABS_ARCH="aarch64"
        NWS_ARCH="aarch64"
        ;;
    *)       
        echo -e "${RED}Unsupported architecture: $ARCH${PLAIN}"
        exit 1 
        ;;
esac

# Download Ookla Speedtest
echo -e "${YELLOW}Downloading Speedtest CLI...${PLAIN}"
SPEEDTEST_URL="https://install.speedtest.net/app/cli/ookla-speedtest-1.2.0-linux-${NWS_ARCH}.tgz"
curl -sL "$SPEEDTEST_URL" -o "$TEMP_DIR/speedtest.tgz"
tar xzf "$TEMP_DIR/speedtest.tgz" -C "$TEMP_DIR"
chmod +x "$TEMP_DIR/speedtest"
SPEEDTEST_CMD="$TEMP_DIR/speedtest"

# Download Iperf3
echo -e "${YELLOW}Downloading Iperf3...${PLAIN}"
IPERF_URL="https://raw.githubusercontent.com/masonr/yet-another-bench-script/master/bin/iperf/iperf3_$YABS_ARCH"
curl -sL "$IPERF_URL" -o "$TEMP_DIR/iperf3"
chmod +x "$TEMP_DIR/iperf3"
IPERF_CMD="$TEMP_DIR/iperf3"

# Initialize JSON output parts
JSON_SPEEDTEST=""
JSON_ROUTING=""
JSON_IPERF=""

# --- 2. Global Speedtest (Smart Resume) ---

echo -e "\n${YELLOW}Running Global Speedtest (nws.sh nodes)...${PLAIN}"

SPEEDTEST_NODES=(
    # --- 欧洲 (6个：覆盖西欧/中欧/东欧/南欧) ---
    "31183,30376,35086|London, UK"      # 伦敦 (西欧核心，移除曼彻斯特)
    "23095,43232,53634|Frankfurt, DE"   # 法兰克福 (欧洲最重要交换中心，移除柏林)
    "37363,22753|Amsterdam, NL"         # 阿姆斯特丹 (西欧另一核心枢纽)
    "61933,26312|Paris, FR"             # 巴黎 (西欧/南欧过渡，移除马赛)
    "7839,11422,3863|Milan, IT"         # 米兰 (南欧核心，移除马德里)
    "14979,10331|Madrid, ES"                # 马德里 (Orange / Movistar)
    "1907,48192,10987,44806|Moscow, RU"             # 莫斯科 (东欧/俄罗斯核心)

    # --- 亚洲 (8个：覆盖东亚/东南亚/南亚亚区) ---
    "68983,13538,1536|Hong Kong, CN"    # 香港 (核心枢纽)
    "18445,2468,3415|Taipei, TW"        # 台北 (东亚重要节点)
    "50686,21507,48463|Tokyo, JP"       # 东京 (东亚核心枢纽)
    "5935,13623,51222|Singapore, SG"    # 新加坡 (东南亚核心枢纽)
    "8990,53900|Bangkok, TH"            # 曼谷 (中南半岛代表，移除河内)
    "13039,26665|Jakarta, ID"           # 雅加达 (东南亚人口大国，保留以测试特定路由)
    "23647,19266|Mumbai, IN"            # 孟买 (南亚核心)
    "2552,36787,35787|Hanoi, VN"            # 河内 (Vinaphone / Viettel / VNPT)

    # --- 北美洲 (7个：覆盖美西/美中/美东/加拿大) ---
    "34840,1782,50081|Los Angeles, US"  # 洛杉矶 (美西核心，CN2/海缆登陆点，移除西雅图)
    "49365,24935|San Jose, US"          # 圣何塞 (硅谷，美西数据中心密集区)
    "22288,16972|Dallas, US"            # 达拉斯 (美中核心，移除丹佛)
    "34707,24056|Chicago, US"           # 芝加哥 (美中北部核心)
    "46120,55152,15647|New York, US"    # 纽约 (美东金融/商业核心)
    "32493,17391|Ashburn, US"           # 阿什本 (美东数据中心核心，全球流量最大节点，移除迈阿密)
    "53393,17300|Toronto, CA"           # 多伦多 (加拿大核心，移除温哥华以节省配额)
    "54754,3499,7945|MX"           # 墨西哥
    
    # --- 中国大陆 (4个：覆盖北/东地理区及三网核心) ---
    "24447|Shanghai, CU-CN"       # 联通 - 上海
    "43752|Beijing, CU-CN"        # 联通 - 北京
    "5396,36663|Suzhou,Zhenjiang, CT-CN"      # 电信 - 苏州镇江
    "59386,59387|Zhejiang, CT-CN"       # 电信 - 浙江
    "16204|Suzhou, CM-CN"         # 移动 - 苏州

    # --- 其他地区 (4个：中东/非洲/南美/大洋洲) ---
    "22129,1692,34240|AE"              # 迪拜 (中东核心，移除吉达/伊斯坦布尔)
    "2962,1491,21570|ZA"      # 约翰内斯堡 (非洲核心)
    "21836,3065,21568|BR"         # 圣保罗 (南美核心)
    "44735,12492,15136,43024|AU"            # 悉尼 (大洋洲核心)
)

count=0
for node in "${SPEEDTEST_NODES[@]}"; do
    ALL_IDS="${node%%|*}"
    NAME="${node##*|}"
    
    printf "%-30s " "Testing $NAME..."

    # --- CHECK JSON FOR EXISTING SUCCESS ---
    # Look for matching location AND download_mbps > 0
    OLD_ENTRY=$(echo "$PREV_JSON" | jq -c --arg loc "$NAME" '.speedtest[]? | select(.location == $loc and (.download_mbps | tonumber > 0))' | head -n 1)

    if [[ -n "$OLD_ENTRY" ]]; then
        echo -e "${CYAN}Skipping (Found valid result)${PLAIN}"
        [[ $count -gt 0 ]] && JSON_SPEEDTEST+=","
        JSON_SPEEDTEST+="$OLD_ENTRY"
        ((count++))
        continue
    fi
    # ---------------------------------------
    
    IFS=',' read -r -a ID_ARRAY <<< "$ALL_IDS"
    RESULT_FOUND=false
    
    for ID in "${ID_ARRAY[@]}"; do
        RESULT=$($SPEEDTEST_CMD -s $ID -f json --accept-license --accept-gdpr 2>/dev/null)
        
        if [[ -n "$RESULT" && "$RESULT" != *"error"* ]]; then
            DL=$(echo "$RESULT" | grep -o '"bandwidth":[0-9]*' | head -1 | awk -F: '{print $2}')
            UL=$(echo "$RESULT" | grep -o '"bandwidth":[0-9]*' | tail -1 | awk -F: '{print $2}')
            LAT=$(echo "$RESULT" | grep -o '"latency":[0-9.]*' | awk -F: '{print $2}')
            
            if [[ -n "$DL" ]]; then DL=$(awk "BEGIN {printf \"%.2f\", $DL * 8 / 1000000}"); fi
            if [[ -n "$UL" ]]; then UL=$(awk "BEGIN {printf \"%.2f\", $UL * 8 / 1000000}"); fi
            
            echo -e "${GREEN}DL: ${DL} Mbps | UL: ${UL} Mbps | Ping: ${LAT} ms${PLAIN}"
            
            [[ $count -gt 0 ]] && JSON_SPEEDTEST+=","
            JSON_SPEEDTEST+="{ \"id\": \"$ID\", \"location\": \"$NAME\", \"download_mbps\": ${DL:-0}, \"upload_mbps\": ${UL:-0}, \"latency_ms\": ${LAT:-0} }"
            
            RESULT_FOUND=true
            break 
        fi
    done
    
    if [[ "$RESULT_FOUND" == "false" ]]; then
        echo -e "${RED}All Nodes Failed${PLAIN}"
        [[ $count -gt 0 ]] && JSON_SPEEDTEST+=","
        JSON_SPEEDTEST+="{ \"id\": \"$ALL_IDS\", \"location\": \"$NAME\", \"download_mbps\": 0, \"upload_mbps\": 0, \"latency_ms\": 0, \"status\": \"failed\" }"
    fi

    ((count++))
done


# --- 3. Global Ping/Routing (Smart Resume) ---

echo -e "\n${YELLOW}Running Global Routing/Ping Test (nws.sh nodes)...${PLAIN}"

ROUTING_NODES=(
    "103.214.168.128|JP - Tokyo: SoftBank"
    "core2.tyo1.he.net|JP - Tokyo: HE"
    "154.18.170.27|HK - Hong Kong: Cogent"
    "43.249.36.49|HK - Hong Kong: LeaseWeb Asia"
    "103.186.65.98|VN - Ho Chi Minh: FPT Telecom"
    "89.187.162.1|SG - Singapore: CDN77"
    "gin-asina-tcore1.as6453.net|SG - Singapore: TATA"
    "speedtest-1.maa1.in.prod.linode.com|IN - Chennai: Linode"
    "49.44.93.128|IN - Mumbai: Jio"
    "37.32.0.1|IR - Tehran: ArvanCloud"
    "194.5.50.94|IR - Isfahan: Webdade"
    "45.142.246.177|RU - Moscow: Misaka"
    "185.45.12.22|RO - Bucharest: M247"
    "fsn1-speed.hetzner.com|DE - Falkenstein: Hetzner"
    "185.94.189.162|FR - Paris: M247"
    "194.127.172.33|NL - Amsterdam: Clouvider"
    "185.59.221.51|UK - London: CDN77"
    "51.222.154.207|CA - Montreal: OVH"
    "192.3.180.103|US - Buffalo: ColoCrossing"
    "104.243.42.233|US - New York: ReliableSite"
    "104.238.204.68|US - Miami: ReliableSite"
    "108.181.140.235|US - Chicago: Psychz Networks"
    "216.146.25.35|US - Dallas: GSL"
    "23.145.48.48|US - Seattle: xTom"
    "64.185.232.162|US - Los Angeles: WebNX"
)

count=0
for node in "${ROUTING_NODES[@]}"; do
    IP="${node%%|*}"
    NAME="${node##*|}"
    
    printf "%-35s " "$NAME"

    # --- CHECK JSON FOR EXISTING SUCCESS ---
    # Look for matching IP AND loss_pct != "100%" (or check ping_avg_ms != N/A)
    OLD_ENTRY=$(echo "$PREV_JSON" | jq -c --arg ip "$IP" '.routing[]? | select(.ip == $ip and .loss_pct != "100%")' | head -n 1)

    if [[ -n "$OLD_ENTRY" ]]; then
        # Extract visual data for printing
        PING_SHOW=$(echo "$OLD_ENTRY" | jq -r .ping_avg_ms)
        LOSS_SHOW=$(echo "$OLD_ENTRY" | jq -r .loss_pct)
        echo -e "${CYAN}${PING_SHOW} ms (${LOSS_SHOW}% loss) [Skipped]${PLAIN}"
        
        [[ $count -gt 0 ]] && JSON_ROUTING+=","
        JSON_ROUTING+="$OLD_ENTRY"
        ((count++))
        continue
    fi
    # ---------------------------------------
    
    # Ping Test (5 packets)
    PING_OUT=$(ping -c 5 -W 2 "$IP" 2>/dev/null)
    if [[ $? -ne 0 ]]; then
        PING_AVG="N/A"
        LOSS="100%"
        echo -e "${RED}Failed${PLAIN}"
    else
        PING_AVG=$(echo "$PING_OUT" | grep -oP 'min/avg/max/mdev = [0-9.]*/\K[0-9.]+')
        LOSS=$(echo "$PING_OUT" | grep -oP '\d+(?=% packet loss)')
        echo -e "${GREEN}${PING_AVG} ms (${LOSS}% loss)${PLAIN}"
    fi

    AS_PATH="N/A"
    if command -v mtr &> /dev/null; then
        AS_PATH=$(timeout 10 mtr -wz -c 1 "$IP" 2>/dev/null | grep "AS[0-9]" | awk '{
            for (i=1; i<=NF; i++) {
                if ($i ~ /^AS/) {
                    if (!seen[$i]++) {
                        printf "%s ", $i
                    }
                }
            }
        }' | sed 's/ $//')
    fi

    [[ $count -gt 0 ]] && JSON_ROUTING+=","
    JSON_ROUTING+="{ \"location\": \"$NAME\", \"ip\": \"$IP\", \"ping_avg_ms\": \"$PING_AVG\", \"loss_pct\": \"$LOSS\", \"as_path\": \"$AS_PATH\" }"
    ((count++))
done


# --- 4. Iperf3 (Smart Resume) ---

echo -e "\n${YELLOW}Running Iperf3 Tests (YABS nodes)...${PLAIN}"

IPERF_NODES=(
    "lon.speedtest.clouvider.net|5200-5209|Clouvider|London, UK (10G)|IPv4|IPv6"
    "iperf-ams-nl.eranium.net|5201-5210|Eranium|Amsterdam, NL (100G)|IPv4|IPv6"
    "speedtest.sin1.sg.leaseweb.net|5201-5210|Leaseweb|Singapore, SG (10G)|IPv4|IPv6"
    "la.speedtest.clouvider.net|5200-5209|Clouvider|Los Angeles, CA, US (10G)|IPv4|IPv6"
    "speedtest.nyc1.us.leaseweb.net|5201-5210|Leaseweb|NYC, NY, US (10G)|IPv4|IPv6"
    "speedtest.sao1.edgoo.net|9204-9240|Edgoo|Sao Paulo, BR (1G)|IPv4|IPv6"
)

run_iperf() {
    local MODE=$1
    local HOST=$2
    local PORTS=$3
    local PROVIDER=$4
    local LOCATION=$5
    
    printf "%-5s | %-15s | %-25s | " "$MODE" "$PROVIDER" "$LOCATION"

    # --- CHECK JSON FOR EXISTING SUCCESS ---
    # Need to match Location AND Mode AND check valid speed
    OLD_ENTRY=$(echo "$PREV_JSON" | jq -c --arg loc "$LOCATION" --arg mode "$MODE" '.iperf[]? | select(.location == $loc and .mode == $mode and (.send_mbps | tonumber > 0))' | head -n 1)

    if [[ -n "$OLD_ENTRY" ]]; then
        S_SHOW=$(echo "$OLD_ENTRY" | jq -r .send_mbps)
        R_SHOW=$(echo "$OLD_ENTRY" | jq -r .recv_mbps)
        echo -e "${CYAN}Up: ${S_SHOW} / Down: ${R_SHOW} [Skipped]${PLAIN}"
        
        [[ $count -gt 0 ]] && JSON_IPERF+=","
        JSON_IPERF+="$OLD_ENTRY"
        ((count++))
        return
    fi
    # ---------------------------------------

    # Pick random port
    local PORT_START=${PORTS%-*}
    local PORT_END=${PORTS#*-}
    local PORT=$(( PORT_START + RANDOM % (PORT_END - PORT_START + 1) ))
    
    local FLAGS=""
    [[ "$MODE" == "IPv6" ]] && FLAGS="-6" || FLAGS="-4"
    
    # SEND TEST
    OUTPUT_SEND=$(timeout 20 $IPERF_CMD -c "$HOST" -p "$PORT" $FLAGS -P 8 -J 2>/dev/null)
    # RECV TEST
    OUTPUT_RECV=$(timeout 20 $IPERF_CMD -c "$HOST" -p "$PORT" $FLAGS -P 8 -J -R 2>/dev/null)
    
    SEND_BPS=$(echo "$OUTPUT_SEND" | grep -oP '"bits_per_second":\s*\K[0-9.]+' | tail -1)
    RECV_BPS=$(echo "$OUTPUT_RECV" | grep -oP '"bits_per_second":\s*\K[0-9.]+' | tail -1)
    
    if [[ -z "$SEND_BPS" ]]; then SEND_VAL="0"; else SEND_VAL=$(awk "BEGIN {printf \"%.2f\", $SEND_BPS / 1000000}"); fi
    if [[ -z "$RECV_BPS" ]]; then RECV_VAL="0"; else RECV_VAL=$(awk "BEGIN {printf \"%.2f\", $RECV_BPS / 1000000}"); fi
    
    if [[ "$SEND_VAL" == "0" && "$RECV_VAL" == "0" ]]; then
         echo -e "${RED}Busy/Fail${PLAIN}"
    else
         echo -e "${GREEN}Up: ${SEND_VAL} Mbps / Down: ${RECV_VAL} Mbps${PLAIN}"
    fi

    [[ $count -gt 0 ]] && JSON_IPERF+=","
    JSON_IPERF+="{ \"mode\": \"$MODE\", \"provider\": \"$PROVIDER\", \"location\": \"$LOCATION\", \"send_mbps\": ${SEND_VAL:-0}, \"recv_mbps\": ${RECV_VAL:-0} }"
    ((count++))
}

count=0
for node in "${IPERF_NODES[@]}"; do
    HOST="${node%%|*}"; R=${node#*|}; PORTS="${R%%|*}"; R=${R#*|}; PROVIDER="${R%%|*}"; R=${R#*|}; LOCATION="${R%%|*}"; MODES="${node##*|}"
    
    if [[ "$HOST_IPV4" != "N/A" ]]; then
        run_iperf "IPv4" "$HOST" "$PORTS" "$PROVIDER" "$LOCATION"
    fi
    if [[ "$HOST_IPV6" != "N/A" && "$MODES" == *"IPv6"* ]]; then
        run_iperf "IPv6" "$HOST" "$PORTS" "$PROVIDER" "$LOCATION"
    fi
done

# --- 5. Generate Final JSON ---

echo -e "\n${YELLOW}Generating JSON Report ($JSON_FILE)...${PLAIN}"

cat <<EOF > "$JSON_FILE"
{
  "host_info": {
    "ipv4": "$HOST_IPV4",
    "ipv6": "$HOST_IPV6",
    "timestamp": "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  },
  "speedtest": [
    $JSON_SPEEDTEST
  ],
  "routing": [
    $JSON_ROUTING
  ],
  "iperf": [
    $JSON_IPERF
  ]
}
EOF

echo -e "${GREEN}Done! Saved to $JSON_FILE${PLAIN}"