#!/usr/bin/env bash
# ==============================================================================
# VH Benchmark Script - VietHosting
# Website: https://viethosting.com
# Version: 1.1
# ==============================================================================

export LC_ALL=C

SCRIPT_VERSION="1.1"
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

LOG_FILE="vh-benchmark-$(date +%Y%m%d-%H%M%S).log"
TMP_PREFIX="/tmp/vhbench_$$"

if ! command -v curl >/dev/null 2>&1; then
    echo -e " ${RED}Error: 'curl' is missing. Please install curl first.${NC}"
    exit 1
fi

cleanup() {
    rm -f "${TMP_PREFIX}"* iotest_$$ fio_test 2>/dev/null
}
trap cleanup INT TERM EXIT

print_separator() {
    echo "  -----------------------------------------------------------------------------------"
}

show_help() {
    echo -e "  ${CYAN}${BOLD}===================================================================================${NC}"
    echo -e "  ${BOLD}Usage Curl:${NC}"
    echo -e "    curl -Lso- https://mirrors.viethosting.com/scripts/vh-benchmark.sh | bash -s [option]\n"
    
    echo -e "  ${BOLD}Usage File:${NC}"
    echo -e "    wget -qO vh-benchmark.sh https://mirrors.viethosting.com/scripts/vh-benchmark.sh"
    echo -e "    bash vh-benchmark.sh [option]\n"
    
    echo -e "  ${BOLD}Options :${NC}"
    printf "    ${GREEN}%-5s${NC} : %s\n" "all" "Run all tests (Default)"
    printf "    ${YELLOW}%-5s${NC} : %s\n" "vn" "Test Vietnam nodes only"
    printf "    ${CYAN}%-5s${NC} : %s\n" "intl" "Test International nodes only"
    printf "    ${BOLD}%-5s${NC} : %s\n" "help" "Show this help menu"
    
    echo -e "  ${CYAN}${BOLD}===================================================================================${NC}"
    echo -e "  ${CYAN}Note:${NC} Network tests include major Vietnamese ISP backbones (VNPT, Viettel, FPT)."
    echo -e "  ${CYAN}Tip:${NC} Advanced users may uncomment additional ISP and routing nodes in the script."
    echo -e "  -----------------------------------------------------------------------------------"
}

main_benchmark() {
    local OPTION=$1
    local START_TIME=$(date +%s)
    
    echo -e "  ${CYAN}${BOLD}===================================================================================${NC}"
    echo -e " ${CYAN}${BOLD}                 VH BENCHMARK SCRIPT v${SCRIPT_VERSION} (https://viethosting.com)   ${NC}"
    echo -e "  ${CYAN}${BOLD}===================================================================================${NC}"

    # ------------------------------------------------------------------------------
    # 1. THÔNG TIN HỆ THỐNG & MẠNG
    # ------------------------------------------------------------------------------
    echo -e "\n  ${YELLOW}[System & Network Information]${NC}"
    print_separator
    
    CPU_MODEL=$(awk -F: '/model name/ {print $2; exit}' /proc/cpuinfo | sed 's/^[ \t]*//')
    CPU_CORES=$(nproc 2>/dev/null || echo "1")
    
    CPU_FREQ=$(lscpu 2>/dev/null | awk -F: '/CPU max MHz|CPU MHz/ {print $2; exit}' | xargs)
    [ -z "$CPU_FREQ" ] && CPU_FREQ=$(awk -F: '/[Cc]pu MHz/ {print $2; exit}' /proc/cpuinfo | xargs)
    [ -z "$CPU_FREQ" ] && CPU_FREQ="N/A"
    
    if command -v free >/dev/null 2>&1 && free -m >/dev/null 2>&1; then
        TOTAL_MEM=$(free -m | awk '/Mem:/ {print $2 "MB (" $3 "MB Used)"}')
        TOTAL_SWAP=$(free -m | awk '/Swap:/ {print $2 "MB (" $3 "MB Used)"}')
    else
        TOTAL_MEM=$(awk '/MemTotal/ {print int($2/1024) "MB"}' /proc/meminfo)
        TOTAL_SWAP=$(awk '/SwapTotal/ {print int($2/1024) "MB"}' /proc/meminfo)
    fi
    
    DISK_INFO=$(df -hT --output=size,used,fstype / 2>/dev/null | awk 'NR==2 {print $1 " (" $2 " Used - Type: " $3 ")"}')
    [ -z "$DISK_INFO" ] && DISK_INFO=$(df -hT / | awk 'NR==2 {print $3 " (" $4 " Used - Type: " $2 ")"}')
    
    PUBLIC_IP=$(curl -sS4 -m 5 ipinfo.io/ip 2>/dev/null)
    if [[ "$PUBLIC_IP" == *"{"* ]] || [[ "$PUBLIC_IP" == *"html"* ]] || [ -z "$PUBLIC_IP" ]; then
        PUBLIC_IP=$(curl -sS4 -m 5 api.ipify.org 2>/dev/null || curl -sS4 -m 5 ifconfig.me 2>/dev/null || echo "Unknown")
    fi

    ASN_ORG=$(curl -sS4 -m 5 ipinfo.io/org 2>/dev/null)
    if [[ "$ASN_ORG" == *"{"* ]] || [[ "$ASN_ORG" == *"html"* ]] || [ -z "$ASN_ORG" ]; then
        ASN_ORG=$(curl -sS4 -m 5 ip-api.com/csv?fields=as,isp 2>/dev/null | tr ',' ' ' || echo "Unknown")
    fi

    OS_NAME=$(grep PRETTY_NAME /etc/os-release 2>/dev/null | cut -d= -f2 | tr -d '"' | sed 's/ *([^)]*)//g')
    KERNEL=$(uname -r)
    ARCH=$(uname -m)
    UPTIME=$(uptime -p 2>/dev/null || uptime | sed 's/.*up //;s/,.*//')
    
    if command -v systemd-detect-virt >/dev/null 2>&1; then 
        VIRT=$(systemd-detect-virt)
    elif command -v virt-what >/dev/null 2>&1; then 
        VIRT=$(virt-what | head -n 1)
    elif grep -qa hypervisor /proc/cpuinfo; then
        VIRT="VM"
    else 
        VIRT="Baremetal"
    fi

    printf "   %-22s: %s\n" "CPU Model" "$CPU_MODEL"
    printf "   %-22s: %s\n" "CPU Cores" "$CPU_CORES Cores @ ${CPU_FREQ} MHz"
    printf "   %-22s: %s\n" "Total RAM" "$TOTAL_MEM"
    printf "   %-22s: %s\n" "Total Swap" "$TOTAL_SWAP"
    printf "   %-22s: %s\n" "Disk Space" "$DISK_INFO"
    printf "   %-22s: %s\n" "OS" "$OS_NAME ($ARCH)"
    printf "   %-22s: %s\n" "Kernel" "$KERNEL"
    printf "   %-22s: %s\n" "Virtualization" "$VIRT"
    printf "   %-22s: %s\n" "Public IPv4" "$PUBLIC_IP"
    printf "   %-22s: %s\n" "ASN / ISP" "$ASN_ORG"	
    printf "   %-22s: %s\n" "Uptime" "$UPTIME"	

    # ------------------------------------------------------------------------------
    # 2. TEST Ổ CỨNG & IO
    # ------------------------------------------------------------------------------
    echo -e "\n  ${YELLOW}[Disk Performance Test]${NC}"
    print_separator
    
    local IO_RUN1="Error" IO_RUN2="Error" IO_RUN3="Error" IO_AVG="0.0"
    local IO_SUM=0
    for i in 1 2 3; do
        local IO_OUTPUT=$(dd if=/dev/zero of=iotest_$$ bs=1M count=1024 oflag=direct 2>&1)
        local IO_SPEED=$(echo "$IO_OUTPUT" | grep -oE '[0-9.]+ ?[KMGTP]i?B/s' | tail -n 1)
        
        local RAW_VAL=$(echo "$IO_SPEED" | tr -d 'a-zA-Z/ ')
        local RAW_UNIT=$(echo "$IO_SPEED" | tr -d '0-9. ')
        
        if [ -z "$RAW_VAL" ]; then RAW_VAL=0; RAW_UNIT="MB/s"; IO_RES="Error"
        else IO_RES="$RAW_VAL $RAW_UNIT"; fi
        
        if [ "$i" -eq 1 ]; then IO_RUN1="$IO_RES"; fi
        if [ "$i" -eq 2 ]; then IO_RUN2="$IO_RES"; fi
        if [ "$i" -eq 3 ]; then IO_RUN3="$IO_RES"; fi
        
        [[ "$RAW_UNIT" == *"G"* ]] && RAW_VAL=$(awk "BEGIN {print $RAW_VAL * 1024}")
        IO_SUM=$(awk "BEGIN {print $IO_SUM + $RAW_VAL}")
        rm -f iotest_$$
    done
    IO_AVG=$(awk "BEGIN {printf \"%.1f\", $IO_SUM / 3}")

    local BW_READ="N/A" IOPS_READ="N/A" BW_WRITE="N/A" IOPS_WRITE="N/A"
    if command -v fio >/dev/null 2>&1; then
        local RES_FIO="${TMP_PREFIX}_fio.log"
        fio --randrepeat=1 --ioengine=libaio --direct=1 --gtod_reduce=1 --name=fio_test --filename=fio_test --bs=4k --numjobs=1 --iodepth=64 --size=256M --readwrite=randrw --rwmixread=75 --runtime=30 --time_based --output="$RES_FIO" > /dev/null 2>&1
        
        local FIO_MAJOR=$(fio -v 2>&1 | grep -o '[0-9]\+' | head -n1)
        if [ "$FIO_MAJOR" = "2" ]; then
            IOPS_READ=$(grep -m1 -i "^ *read *:" "$RES_FIO" | awk -F'[=, ]+' '{for(i=1;i<=NF;i++) if($i=="iops") print $(i+1)}')
            BW_READ=$(grep -m1 -i "^ *read *:" "$RES_FIO" | awk -F'[=, ]+' '{for(i=1;i<=NF;i++) if($i=="bw") print $(i+1)}')
            IOPS_WRITE=$(grep -m1 -i "^ *write *:" "$RES_FIO" | awk -F'[=, ]+' '{for(i=1;i<=NF;i++) if($i=="iops") print $(i+1)}')
            BW_WRITE=$(grep -m1 -i "^ *write *:" "$RES_FIO" | awk -F'[=, ]+' '{for(i=1;i<=NF;i++) if($i=="bw") print $(i+1)}')
        else
            IOPS_READ=$(grep -m1 -E -i "^ *read *:" "$RES_FIO" | grep -oEi 'iops=[0-9.]+[a-z]*' | cut -d= -f2)
            BW_READ=$(grep -m1 -E -i "^ *read *:" "$RES_FIO" | grep -oEi 'bw=[0-9.]+[a-z/]+' | cut -d= -f2 | sed -e 's/MiB/MB/gi' -e 's/KiB/KB/gi')
            IOPS_WRITE=$(grep -m1 -E -i "^ *write *:" "$RES_FIO" | grep -oEi 'iops=[0-9.]+[a-z]*' | cut -d= -f2)
            BW_WRITE=$(grep -m1 -E -i "^ *write *:" "$RES_FIO" | grep -oEi 'bw=[0-9.]+[a-z/]+' | cut -d= -f2 | sed -e 's/MiB/MB/gi' -e 's/KiB/KB/gi')
        fi
        [ -z "$IOPS_READ" ] && IOPS_READ="N/A"
        [ -z "$BW_READ" ] && BW_READ="N/A"
        [ -z "$IOPS_WRITE" ] && IOPS_WRITE="N/A"
        [ -z "$BW_WRITE" ] && BW_WRITE="N/A"
        rm -f "$RES_FIO" fio_test
    else
        BW_READ="${CYAN}Missing 'fio'${NC}"; IOPS_READ="${CYAN}Missing 'fio'${NC}"
        BW_WRITE="${CYAN}Missing 'fio'${NC}"; IOPS_WRITE="${CYAN}Missing 'fio'${NC}"
    fi

    echo -e "   ${BOLD}Sequential Disk Test (dd - 3 times)${NC} |   ${BOLD}Random 4K Read/Write (fio)${NC}"
	printf "   %-22s: %-11s |   %-22s: %b\n" "Test 1" "$IO_RUN1" "Fio Read (75%)" "$BW_READ"
    printf "   %-22s: %-11s |   %-22s: %b\n" "Test 2" "$IO_RUN2" "Read IOPS" "$IOPS_READ"
    printf "   %-22s: %-11s |   %-22s: %b\n" "Test 3" "$IO_RUN3" "Fio Write (25%)" "$BW_WRITE"
    printf "   %-22s: %-11s |   %-22s: %b\n" "Average" "${IO_AVG} MB/s" "Write IOPS" "$IOPS_WRITE"

    # ------------------------------------------------------------------------------
    # 3. TEST NETWORK 
    # ------------------------------------------------------------------------------
    speed_test_node() {
        local NODE_NAME="$1"; local URL="$2"
        local CURL_OUT=$(curl -A "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" --connect-timeout 8 -m 30 -4 -Lso /dev/null -w "%{remote_ip}|%{speed_download}" "$URL" 2>/dev/null)
        local IP=$(echo "$CURL_OUT" | cut -d'|' -f1)
        local SPEED_BPS=$(echo "$CURL_OUT" | cut -d'|' -f2)
        
        if [ -z "$IP" ] || [ -z "$SPEED_BPS" ] || [ "$SPEED_BPS" == "0.000" ]; then
            printf " %-35s | %-16s | ${RED}%-12s${NC} | %s\n" "$NODE_NAME" "Timeout/Fail" "N/A" "N/A"
            return
        fi
        
        read SPEED_MBPS SPEED_MBIT <<< $(awk -v s="$SPEED_BPS" 'BEGIN {printf "%.2f %.2f", s/1024/1024, (s/1024/1024)*8}')
        local COLOR=$(awk -v speed="$SPEED_MBPS" -v red="$RED" -v yellow="$YELLOW" -v green="$GREEN" 'BEGIN { if (speed < 10) print red; else if (speed < 50) print yellow; else print green; }')
        
        printf "   %-35s | %-16s | ${COLOR}%-12s${NC} | ${CYAN}%s Mbps${NC}\n" "$NODE_NAME" "$IP" "$SPEED_MBPS MB/s" "$SPEED_MBIT"
    }

    run_vn_nodes() {
        echo -e "\n  ${YELLOW}[Vietnam Network Nodes - IPv4]${NC}"
        print_separator
        printf "   ${BOLD}%-35s | %-16s | %-12s | %s${NC}\n" "Node Name" "IPv4 Address" "Download" "Bandwidth"
        print_separator
        speed_test_node "VNPT, Ha Noi, VN" "https://vnpt-hn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        #speed_test_node "VNPT, Da Nang, VN" "https://vnpt-dn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        speed_test_node "VNPT, Ho Chi Minh, VN" "https://vnpt-hcmc01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        #speed_test_node "Viettel, Ha Noi, VN" "https://viettel-hn03.i-speed.vn/garbage.php?cors=true&ckSize=100"
		speed_test_node "Viettel, Ha Noi, VN" "https://viettel-hn04.i-speed.vn/garbage.php?cors=true&ckSize=100"
        #speed_test_node "Viettel, Da Nang, VN" "https://viettel-dn02.i-speed.vn/garbage.php?cors=true&ckSize=100"
        speed_test_node "Viettel, Ho Chi Minh, VN" "https://viettel-hcmc04.i-speed.vn/garbage.php?cors=true&ckSize=100"
		#speed_test_node "Viettel, Ho Chi Minh, VN" "https://viettel-hcmc03.i-speed.vn/garbage.php?cors=true&ckSize=100"
		speed_test_node "Viettel, Binh Duong, VN" "https://viettelidc-bd01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        speed_test_node "FPT Telecom, Ha Noi, VN" "https://fpt-hn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        speed_test_node "FPT Telecom, Ho Chi Minh, VN" "https://fpt-hcmc01.i-speed.vn/garbage.php?cors=true&ckSize=100"
		
		# ------------------
        # [ ADVANCED USERS ] Uncomment the lines below to test specific VN ISPs
        # ------------------		
        
        # speed_test_node "Mobifone, Ha Noi, VN" "https://mobifone-hn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
		# speed_test_node "Mobifone 02, Ha Noi, VN" "https://mobifone-hn02.i-speed.vn/garbage.php?cors=true&ckSize=100"
        # speed_test_node "Mobifone, Da Nang, VN" "https://mobifone-dn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
		# speed_test_node "Mobifone 02, Da Nang, VN" "https://mobifone-dn02.i-speed.vn/garbage.php?cors=true&ckSize=100"
        # speed_test_node "Mobifone, Ho Chi Minh, VN" "https://mobifone-hcmc01.i-speed.vn/garbage.php?cors=true&ckSize=100"
		# speed_test_node "Mobifone 02, Ho Chi Minh, VN" "https://mobifone-hcmc02.i-speed.vn/garbage.php?cors=true&ckSize=100"
		
        # speed_test_node "CMC Telecom, Ha Noi, VN" "https://cmc-hn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        # speed_test_node "CMC Telecom, Da Nang, VN" "https://cmc-dn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        # speed_test_node "CMC Telecom, Ho Chi Minh, VN" "https://cmc-hcmc01.i-speed.vn/garbage.php?cors=true&ckSize=100"
		
        # speed_test_node "SCTV, Ha Noi, VN" "https://sctv-hn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
		# speed_test_node "SCTV, Da Nang, VN" "https://sctv-dn01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        # speed_test_node "SCTV, Ho Chi Minh, VN" "https://sctv-hcmc01.i-speed.vn/garbage.php?cors=true&ckSize=100"		       
    }

    run_intl_nodes() {
        echo -e "\n  ${YELLOW}[International Network Nodes - IPv4]${NC}"
        print_separator
        printf "   ${BOLD}%-35s | %-16s | %-12s | %s${NC}\n" "Node Name" "IPv4 Address" "Download" "Bandwidth"
        print_separator
        
        # --- Global CDNs ---
        speed_test_node "Cloudflare CDN, Global" "https://speed.cloudflare.com/__down?bytes=10000000"
        # Remove cachefly - too small test file (1MB), the result will be incorrect
        # speed_test_node "CacheFly CDN, Global" "https://cachefly.cachefly.net/1mb.test"
        
        # --- Asia ---
        speed_test_node "DataPacket, Hong Kong, HK" "https://hkg.download.datapacket.com/100mb.bin"
        speed_test_node "OVH, Singapore, SG" "https://sgp.proof.ovh.net/files/100Mb.dat"
        speed_test_node "Vultr, Tokyo, JP" "https://hnd-jp-ping.vultr.com/vultr.com.100MB.bin"
        # speed_test_node "Vultr, Seoul, KR" "https://sel-kor-ping.vultr.com/vultr.com.100MB.bin"
        # speed_test_node "OVH, Mumbai, IN" "https://bom.proof.ovh.net/files/100Mb.dat"
        
        # --- Americas ---
        speed_test_node "Vultr, Los Angeles, US West" "https://lax-ca-us-ping.vultr.com/vultr.com.100MB.bin"
        # speed_test_node "Vultr, New York, US East" "https://nj-us-ping.vultr.com/vultr.com.100MB.bin"
        # speed_test_node "Vultr, Chicago, US Central" "https://il-us-ping.vultr.com/vultr.com.100MB.bin"
        speed_test_node "OVH, Montreal, CA" "https://bhs.proof.ovh.ca/files/100Mb.dat"
        # speed_test_node "Vultr, Sao Paulo, BR" "https://sao-br-ping.vultr.com/vultr.com.100MB.bin"

        # --- Europe & Oceania ---
        speed_test_node "DataPacket, Frankfurt, DE" "https://fra.download.datapacket.com/100mb.bin"
        speed_test_node "Vultr, London, UK" "https://lon-gb-ping.vultr.com/vultr.com.100MB.bin"
        # speed_test_node "Vultr, Amsterdam, NL" "https://ams-nl-ping.vultr.com/vultr.com.100MB.bin"
        # speed_test_node "Vultr, Paris, FR" "https://par-fr-ping.vultr.com/vultr.com.100MB.bin"
        speed_test_node "OVH, Sydney, AU" "https://syd.proof.ovh.net/files/100Mb.dat"
		
		# ------------------
        # [ ADVANCED USERS ] Specific Routing 
		# ------------------
		
        # speed_test_node "IPTP Networks, Singapore, SG" "https://iptp-sg01.i-speed.vn/garbage.php?cors=true&ckSize=100"
        # speed_test_node "APNIC, Brisbane, AU" "https://apnic-01.i-speed.vn/garbage.php?cors=true&ckSize=100"
    }

    case "$OPTION" in
        vn) run_vn_nodes ;;
        intl) run_intl_nodes ;;
        all|"") run_vn_nodes; run_intl_nodes ;;
    esac

    local END_TIME=$(date +%s)
    local TOTAL_TIME=$((END_TIME - START_TIME))
    local TIME_MIN=$((TOTAL_TIME / 60))
    local TIME_SEC=$((TOTAL_TIME % 60))
    
	echo -e "\n  ${CYAN}${BOLD}===================================================================================${NC}"
    echo -e "   ${BOLD}Finished in${NC} : ${CYAN}${TIME_MIN} min ${TIME_SEC} sec${NC}. ${BOLD}Log${NC} : ${CYAN}$(pwd)/${LOG_FILE}${NC}"
    echo -e "  ${CYAN}${BOLD}===================================================================================${NC}"
}


case "$1" in
    -h|--help|help)
        show_help
        exit 0
        ;;
    vn|intl|all|"")
        main_benchmark "$1" | tee >(sed 's/\x1b\[[0-9;]*[a-zA-Z]//g' > "$LOG_FILE")
        ;;
    *)
        echo -e "  ${RED}Invalid Option!${NC} Run 'bash vh-benchmark.sh -h' for help."
        exit 1
        ;;
esac

unset TMP_PREFIX SCRIPT_VERSION