#!/bin/bash

set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
WHITE='\033[1;37m'
NC='\033[0m'

declare -A DNS_SERVERS=(
    ["Cloudflare"]="1.1.1.1,1.0.0.1,2606:4700:4700::1111,2606:4700:4700::1001"
    ["Google"]="8.8.8.8,8.8.4.4,2001:4860:4860::8888,2001:4860:4860::8844"
    ["Level3"]="4.2.2.1,4.2.2.2,4.2.2.3,2001:504:8302::53"
    ["OpenDNS"]="208.67.222.222,208.67.220.220,2620:0:ccc::2,2620:0:ccd::2"
    ["Quad9"]="9.9.9.9,,2620:fe::fe"
    ["Comodo"]="8.26.56.26,8.20.247.20,"
)

declare -A TEST_TARGETS=(
    ["browsing"]="google.com,youtube.com,facebook.com,wikipedia.org"
    ["gaming"]="steampowered.com,epicgames.com,104.160.131.3,104.160.141.3,185.40.64.1,192.64.150.1,185.60.112.157"
    ["server"]="github.com,stackoverflow.com,amazon.com,cloudflare.com"
)

declare -A RESULTS=()
declare -A DNS_NAMES=()

show_banner() {
    clear
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${WHITE}                    DNS Performance Tester                    ${CYAN}║${NC}"
    echo -e "${CYAN}║${WHITE}              Professional Network Analysis Tool              ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo
}

show_menu() {
    echo -e "${YELLOW}Select your primary usage goal:${NC}"
    echo -e "${GREEN}1)${NC} Web Browsing & General Use"
    echo -e "${GREEN}2)${NC} Gaming & Low Latency"
    echo -e "${GREEN}3)${NC} Server & Development"
    echo -e "${GREEN}4)${NC} Custom Test Targets"
    echo -e "${GREEN}5)${NC} Exit"
    echo
    read -p "Enter your choice (1-5): " choice
}

get_test_targets() {
    local goal=$1
    case $goal in
        1) echo "${TEST_TARGETS[browsing]}" ;;
        2) echo "${TEST_TARGETS[gaming]}" ;;
        3) echo "${TEST_TARGETS[server]}" ;;
        4) 
            echo "Enter custom targets (comma-separated):"
            read -p "Targets: " custom_targets
            echo "$custom_targets"
            ;;
        *) echo "${TEST_TARGETS[browsing]}" ;;
    esac
}

test_dns_performance() {
    local dns_name=$1
    local dns_servers=$2
    local targets=$3
    
    echo -e "${BLUE}Testing $dns_name DNS...${NC}"
    
    local total_latency=0
    local successful_tests=0
    local test_count=0
    
    # Get the first IPv4 DNS server for testing
    local primary_dns=$(echo "$dns_servers" | cut -d',' -f1)
    
    IFS=',' read -ra TARGET_ARRAY <<< "$targets"
    for target in "${TARGET_ARRAY[@]}"; do
        target=$(echo "$target" | xargs)
        if [[ -n "$target" ]]; then
            test_count=$((test_count + 1))
            local latency=$(measure_dns_query_time "$target" "$primary_dns" 2>/dev/null || echo "999")
            if [[ "$latency" != "999" ]]; then
                total_latency=$((total_latency + latency))
                successful_tests=$((successful_tests + 1))
            fi
        fi
    done
    
    if [[ $successful_tests -gt 0 ]]; then
        local avg_latency=$((total_latency / successful_tests))
        RESULTS["$dns_name"]=$avg_latency
        DNS_NAMES["$dns_name"]="$dns_servers"
    else
        RESULTS["$dns_name"]=999
        DNS_NAMES["$dns_name"]="$dns_servers"
    fi
}

measure_dns_query_time() {
    local target=$1
    local dns_server=$2
    
    local latency=999
    
    # Use dig to measure actual DNS query time
    if command -v dig &> /dev/null; then
        # First verify the query works
        local dig_output=$(timeout 5 dig @"$dns_server" "$target" A +short 2>/dev/null)
        if [[ -n "$dig_output" ]]; then
            # Measure the actual query time
            local start_time=$(date +%s%3N)
            timeout 5 dig @"$dns_server" "$target" A +short >/dev/null 2>&1
            local end_time=$(date +%s%3N)
            latency=$(( (end_time - start_time) ))
            
            # Convert to milliseconds if needed
            if [[ $latency -gt 1000 ]]; then
                latency=$(( latency / 1000 ))
            fi
        fi
    elif command -v nslookup &> /dev/null; then
        # Fallback to nslookup with timing
        local start_time=$(date +%s%3N)
        local nslookup_output=$(timeout 5 nslookup "$target" "$dns_server" 2>/dev/null)
        local end_time=$(date +%s%3N)
        
        if echo "$nslookup_output" | grep -q "Address:"; then
            latency=$(( (end_time - start_time) ))
            
            # Convert to milliseconds if needed
            if [[ $latency -gt 1000 ]]; then
                latency=$(( latency / 1000 ))
            fi
        fi
    fi
    
    # Ensure latency is reasonable (not negative or too high)
    if [[ $latency -lt 0 || $latency -gt 5000 ]]; then
        latency=999
    fi
    
    echo "$latency"
}

run_tests() {
    local targets=$1
    
    echo -e "${YELLOW}Starting DNS performance tests...${NC}"
    echo -e "${CYAN}Test targets: $targets${NC}"
    echo
    
    for dns_name in "${!DNS_SERVERS[@]}"; do
        test_dns_performance "$dns_name" "${DNS_SERVERS[$dns_name]}" "$targets"
    done
}

display_results() {
    echo
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${WHITE}                        TEST RESULTS                          ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo
    
    local sorted_dns=($(for dns in "${!RESULTS[@]}"; do echo "$dns:${RESULTS[$dns]}"; done | sort -t: -k2 -n | cut -d: -f1))
    
    printf "${WHITE}%-12s %-8s %-20s %-20s${NC}\n" "Provider" "Latency" "IPv4 Addresses" "IPv6 Addresses"
    echo -e "${PURPLE}────────────────────────────────────────────────────────────────────────${NC}"
    
    local rank=1
    for dns_name in "${sorted_dns[@]}"; do
        local latency=${RESULTS[$dns_name]}
        local dns_servers=${DNS_NAMES[$dns_name]}
        
        local ipv4_servers=$(echo "$dns_servers" | cut -d',' -f1-2 | tr ',' ' ')
        local ipv6_servers=$(echo "$dns_servers" | cut -d',' -f3-4 | tr ',' ' ')
        
        if [[ "$latency" == "999" ]]; then
            printf "${RED}#%-2d %-10s %-8s %-20s %-20s${NC}\n" "$rank" "$dns_name" "Failed" "$ipv4_servers" "$ipv6_servers"
        else
            local color=""
            if [[ $rank -eq 1 ]]; then
                color="${GREEN}"
            elif [[ $rank -eq 2 ]]; then
                color="${YELLOW}"
            elif [[ $rank -eq 3 ]]; then
                color="${BLUE}"
            else
                color="${WHITE}"
            fi
            printf "${color}#%-2d %-10s %-8s %-20s %-20s${NC}\n" "$rank" "$dns_name" "${latency}ms" "$ipv4_servers" "$ipv6_servers"
        fi
        rank=$((rank + 1))
    done
    
    echo
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${WHITE}                    RECOMMENDED DNS SERVERS                   ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo
    
    local best_dns=${sorted_dns[0]}
    local best_servers=${DNS_NAMES[$best_dns]}
    local best_ipv4=$(echo "$best_servers" | cut -d',' -f1-2)
    local best_ipv6=$(echo "$best_servers" | cut -d',' -f3-4)
    
    echo -e "${GREEN}Best performing DNS: $best_dns${NC}"
    echo
    echo -e "${YELLOW}IPv4 Addresses:${NC}"
    IFS=',' read -ra IPV4_ARRAY <<< "$best_ipv4"
    for ip in "${IPV4_ARRAY[@]}"; do
        if [[ -n "$ip" ]]; then
            echo -e "  ${GREEN}•${NC} $ip"
        fi
    done
    
    echo
    echo -e "${YELLOW}IPv6 Addresses:${NC}"
    IFS=',' read -ra IPV6_ARRAY <<< "$best_ipv6"
    for ip in "${IPV6_ARRAY[@]}"; do
        if [[ -n "$ip" ]]; then
            echo -e "  ${GREEN}•${NC} $ip"
        fi
    done
    
    echo
    echo -e "${BLUE}Configuration Instructions:${NC}"
    echo -e "${WHITE}1.${NC} Open your network settings"
    echo -e "${WHITE}2.${NC} Set DNS servers to the addresses above"
    echo -e "${WHITE}3.${NC} Restart your network connection"
    echo -e "${WHITE}4.${NC} Test your connection speed"
}

show_post_results_menu() {
    echo
    echo -e "${CYAN}╔══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${CYAN}║${WHITE}                         OPTIONS                          ${CYAN}║${NC}"
    echo -e "${CYAN}╚══════════════════════════════════════════════════════════════╝${NC}"
    echo
    echo -e "${GREEN}1)${NC} Restart Test"
    echo -e "${RED}2)${NC} Quit"
    echo
    read -p "Enter your choice (1-2): " post_choice
}

check_dependencies() {
    local missing_deps=()
    
    if ! command -v dig &> /dev/null && ! command -v nslookup &> /dev/null; then
        missing_deps+=("dnsutils or bind-utils")
    fi
    
    if [[ ${#missing_deps[@]} -gt 0 ]]; then
        echo -e "${RED}Missing dependencies: ${missing_deps[*]}${NC}"
        echo -e "${YELLOW}Please install them and run the script again.${NC}"
        echo -e "${CYAN}Example: sudo apt install dnsutils${NC}"
        echo -e "${CYAN}Or: sudo yum install bind-utils${NC}"
        exit 1
    fi
    
    if command -v dig &> /dev/null; then
        echo -e "${GREEN}Using dig for DNS queries (recommended)${NC}"
    else
        echo -e "${YELLOW}Using nslookup for DNS queries (dig preferred)${NC}"
    fi
}

main() {
    check_dependencies
    
    while true; do
        show_banner
        show_menu
        
        case $choice in
            1|2|3|4)
                local targets=$(get_test_targets "$choice")
                if [[ -z "$targets" ]]; then
                    echo -e "${YELLOW}No targets specified. Using default.${NC}"
                    targets="${TEST_TARGETS[browsing]}"
                fi
                
                run_tests "$targets"
                display_results
                
                while true; do
                    show_post_results_menu
                    case $post_choice in
                        1)
                            break
                            ;;
                        2)
                            echo -e "${GREEN}Goodbye!${NC}"
                            exit 0
                            ;;
                        *)
                            echo -e "${RED}Invalid choice. Please try again.${NC}"
                            sleep 1
                            ;;
                    esac
                done
                ;;
            5)
                echo -e "${GREEN}Goodbye!${NC}"
                exit 0
                ;;
            *)
                echo -e "${RED}Invalid choice. Please try again.${NC}"
                sleep 2
                ;;
        esac
    done
}

main "$@"