package dns

import (
	"net"
	"strings"
	"time"

	"github.com/zeevdiukman/test/app/helper"

	miekgDNS "github.com/miekg/dns"
)

// HandleTypeNS handles DNS NS queries.
func (a *DNS) HandleTypeNS(w miekgDNS.ResponseWriter, r *miekgDNS.Msg) {
	t := time.Now() // Start timer
	msg := &miekgDNS.Msg{}
	msg.SetReply(r)                   // Set reply message
	q := r.Question[0]                // Get question
	domainName := helper.FtoD(q.Name) // Convert FQDN to domain name
	ns := ""
	if nsValue, ok := a.Records.TypeNS[domainName]; ok {
		ns = nsValue // Get NS from records
	} else {
		ns = "default-ns.example.com." // Default NS if not found
	}
	RR_Header := miekgDNS.RR_Header{
		Name:   miekgDNS.Fqdn(domainName), // Set domain name
		Rrtype: miekgDNS.TypeNS,           // Set record type
		Class:  miekgDNS.ClassINET,        // Set class
		Ttl:    3600,                      // Set TTL
	}
	answer_typeNS := &miekgDNS.NS{
		Hdr: RR_Header, // Set header
		Ns:  ns,        // Set NS
	}
	msg.Authoritative = true                       // Set authoritative flag
	msg.RecursionDesired = false                   // Set recursion desired flag
	msg.SetRcode(r, miekgDNS.RcodeSuccess)         // Set response code
	msg.Answer = append(msg.Answer, answer_typeNS) // Add answer
	tt := time.Since(t)                            // Calculate elapsed time
	helper.P(tt, " => ", domainName)               // Print elapsed time and domain name

	w.WriteMsg(msg) // Write response
}

// HandleTypeA handles DNS Type A queries.
func (a *DNS) HandleTypeA(w miekgDNS.ResponseWriter, r *miekgDNS.Msg) {
	useAlternateResolver := false
	t := time.Now() // Start timer
	msg := &miekgDNS.Msg{}
	msg.SetReply(r)                   // Set reply message
	q := r.Question[0]                // Get question
	domainName := helper.FtoD(q.Name) // Convert FQDN to domain name
	ip := ""
	if ipValue, ok := a.Records.TypeA[domainName]; ok {
		ip = ipValue // Get IP from records
	} else {
		dSlices := strings.Split(domainName, ".") // Split domain name
		// Check if wildcard
		if len(dSlices) > 2 {
			name := dSlices[len(dSlices)-2]
			tld := dSlices[len(dSlices)-1]
			cname := name + "." + tld
			wildCard := "*." + cname
			if ipValue, ok := a.Records.TypeA[wildCard]; ok {
				ip = ipValue // Get IP from wildcard record
			} else {
				useAlternateResolver = true // Use alternate resolver
			}
		} else {
			useAlternateResolver = true // Use alternate resolver
		}

		if useAlternateResolver {
			ip = a.Resolver.Lookup(domainName) // Lookup IP using resolver
		}
	}
	RR_Header := miekgDNS.RR_Header{
		Name:   miekgDNS.Fqdn(domainName), // Set domain name
		Rrtype: miekgDNS.TypeA,            // Set record type
		Class:  miekgDNS.ClassINET,        // Set class
		Ttl:    3600,                      // Set TTL
	}
	answer_typeA := &miekgDNS.A{
		Hdr: RR_Header,             // Set header
		A:   net.ParseIP(ip).To4(), // Set IP address
	}
	msg.Authoritative = true                      // Set authoritative flag
	msg.RecursionDesired = false                  // Set recursion desired flag
	msg.SetRcode(r, miekgDNS.RcodeSuccess)        // Set response code
	msg.Answer = append(msg.Answer, answer_typeA) // Add answer
	tt := time.Since(t)                           // Calculate elapsed time
	helper.P(tt, " => ", domainName)              // Print elapsed time and domain name

	w.WriteMsg(msg) // Write response
}