package dns

import (
	"fmt"
	"github.com/miekg/dns"
	"net"
	"slices"
	"strings"
)

type Domain struct {
	Name string
	NS   []string
}

func CheckDomain(d, tld string) Domain {
	nameservers := []string{}
	fqdn := fmt.Sprintf("%v.%v", d, tld)

	tldServers, err := getNameserverForTLD(tld)
	if err != nil {
		return Domain{fqdn, nameservers}
	}

	for _, tldServer := range tldServers {
		nameservers, err = queryTLDNamserver(fqdn, tldServer)
		if err == nil {
			break
		}
	}

	if err != nil {
		return Domain{
			Name: fqdn,
			NS:   nameservers,
		}
	}

	slices.Sort(nameservers)
	return Domain{
		Name: fqdn,
		NS:   nameservers,
	}
}

func getNameserverForTLD(tld string) ([]string, error) {
	ns, err := net.LookupNS(tld)
	if err != nil {
		return nil, err
	}
	var tldServer []string
	if len(ns) > 0 {
		for _, n := range ns {
			tldServer = append(tldServer, strings.Trim(n.Host, ".")+":53")
		}
	} else {
		return nil, fmt.Errorf("no nameserver found for TLD %v", tld)
	}
	return tldServer, nil
}

func queryTLDNamserver(fqdn, tldServer string) ([]string, error) {
	nameservers := []string{}
	// Create a new DNS message
	m := new(dns.Msg)
	// Set the question section of the message with the fqdn and query type
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeNS)

	// Create a DNS client
	c := new(dns.Client)

	// Send the DNS query to the TLD server
	r, _, err := c.Exchange(m, tldServer)
	if err != nil {
		return nameservers, fmt.Errorf("dns query error: %v", err)
	}

	// Check for response errors
	if r.Rcode != dns.RcodeSuccess {
		return nameservers, fmt.Errorf("failed to get a valid response for %v: %v", fqdn, r.Rcode)
	}

	// Print the nameservers from the response
	for _, ans := range r.Ns {
		if ns, ok := ans.(*dns.NS); ok {
			nameservers = append(nameservers, strings.ToLower(ns.Ns))
		}
	}
	return nameservers, nil
}