package utils

import (
	"fmt"
	"net"
	"time"

	"github.com/miekg/dns"
)

func ValidateDNSPropagation(hostname, expectedIP string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	backoff := 2 * time.Second
	maxBackoff := 10 * time.Second

	dnsServers := []string{"1.1.1.1:53", "8.8.8.8:53", "1.0.0.1:53"}

	for {
		if time.Now().After(deadline) {
			return fmt.Errorf("timeout waiting for DNS propagation for %s", hostname)
		}

		resolved := false
		
		// Try public DNS servers first
		for _, dnsServer := range dnsServers {
			ips, err := queryDNS(dnsServer, hostname)
			if err == nil {
				for _, ip := range ips {
					if ip == expectedIP {
						resolved = true
						break
					}
				}
				if resolved {
					break
				}
			}
		}

		// Fallback to system resolver if public DNS servers failed
		if !resolved {
			ips, err := net.LookupIP(hostname)
			if err == nil {
				for _, ip := range ips {
					if ip.String() == expectedIP {
						resolved = true
						break
					}
				}
			}
		}

		if resolved {
			return nil
		}

		time.Sleep(backoff)
		if backoff < maxBackoff {
			backoff += 1 * time.Second
		}
	}
}

func queryDNS(server, hostname string) ([]string, error) {
	client := dns.Client{Timeout: 5 * time.Second}
	msg := new(dns.Msg)
	msg.SetQuestion(dns.Fqdn(hostname), dns.TypeA)

	resp, _, err := client.Exchange(msg, server)
	if err != nil {
		return nil, err
	}

	if resp.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("DNS query failed with Rcode: %d", resp.Rcode)
	}

	var ips []string
	for _, answer := range resp.Answer {
		if a, ok := answer.(*dns.A); ok {
			ips = append(ips, a.A.String())
		}
	}

	return ips, nil
}

func ResolveHostname(hostname string) ([]string, error) {
	ips, err := net.LookupIP(hostname)
	if err != nil {
		return nil, err
	}

	var ipStrings []string
	for _, ip := range ips {
		if ip.To4() != nil {
			ipStrings = append(ipStrings, ip.String())
		}
	}

	return ipStrings, nil
}

func ValidateDNSPropagationWithResolver(hostname, expectedIP, resolver string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	backoff := 2 * time.Second
	maxBackoff := 10 * time.Second

	dnsServer := resolver + ":53"

	for {
		if time.Now().After(deadline) {
			return fmt.Errorf("timeout waiting for DNS propagation for %s on resolver %s", hostname, resolver)
		}

		ips, err := queryDNS(dnsServer, hostname)
		if err == nil {
			for _, ip := range ips {
				if ip == expectedIP {
					return nil
				}
			}
		}

		time.Sleep(backoff)
		if backoff < maxBackoff {
			backoff += 1 * time.Second
		}
	}
}