package traverse

import (
	"net"

	"gitea.hansenits.com.au/hits/ExploreDNS/internal/dns"
	miekgdns "github.com/miekg/dns"
)

type ResponseType int

const (
	RespReferral ResponseType = iota
	RespAnswer
	RespCNAMEFollow
	RespNODATA
	RespNXDOMAIN
	RespSERVFAIL
	RespREFUSED
	RespNOTIMPL
	RespCNAMELoop
	RespError
	// RespNSResolutionFailed indicates that the traversal could not resolve the
	// IP address of an in-bailiwick nameserver.  The domain may still be
	// reachable in practice (e.g. via glue records held by the registry), but
	// the iterative traversal could not complete that path.
	RespNSResolutionFailed
)

func (rt ResponseType) String() string {
	switch rt {
	case RespReferral:
		return "referral"
	case RespAnswer:
		return "answer"
	case RespCNAMEFollow:
		return "cname_follow"
	case RespNODATA:
		return "nodata"
	case RespNXDOMAIN:
		return "nxdomain"
	case RespSERVFAIL:
		return "servfail"
	case RespREFUSED:
		return "refused"
	case RespNOTIMPL:
		return "notimp"
	case RespCNAMELoop:
		return "cname_loop"
	case RespError:
		return "error"
	case RespNSResolutionFailed:
		return "ns_error"
	default:
		return "unknown"
	}
}

type Response struct {
	Referral     *Referral
	Server       net.IP
	Cache        *InfoCache
	Decoded      *dns.DecodedResponse
	Type         ResponseType
	ErrorMessage string
}

func NewResponse(ref *Referral, server net.IP, cache *InfoCache) *Response {
	return &Response{
		Referral: ref,
		Server:   server,
		Cache:    cache,
	}
}

func (r *Response) Process(msg *miekgdns.Msg) *Response {
	if msg == nil {
		r.Type = RespError
		r.ErrorMessage = "nil DNS response"
		return r
	}

	r.Decoded = dns.DecodeResponse(msg)
	if r.Decoded == nil {
		r.Type = RespError
		r.ErrorMessage = "failed to decode DNS response"
		return r
	}

	// Synthesize CNAME from DNAME when the server didn't include a synthesized CNAME record.
	if len(r.Decoded.CNAMEChain) == 0 && r.Referral != nil && len(r.Decoded.DNAMEMappings) > 0 {
		for _, dm := range r.Decoded.DNAMEMappings {
			synthesized := dns.SynthesizeCNAMEFromDNAME(r.Referral.Name, dm.Owner, dm.Target)
			if synthesized != "" {
				r.Decoded.CNAMEChain = append(r.Decoded.CNAMEChain, synthesized)
				break
			}
		}
	}

	r.Type = r.classify()
	return r
}

func (r *Response) classify() ResponseType {
	switch r.Decoded.Classification {
	case dns.ResponseNXDOMAIN:
		return RespNXDOMAIN
	case dns.ResponseSERVFAIL:
		return RespSERVFAIL
	case dns.ResponseREFUSED:
		return RespREFUSED
	case dns.ResponseNOTIMPL:
		return RespNOTIMPL
	case dns.ResponseAnswer:
		if len(r.Decoded.CNAMEChain) > 0 && !r.hasFinalAnswer() {
			return RespCNAMEFollow
		}
		return RespAnswer
	case dns.ResponseReferral:
		return RespReferral
	case dns.ResponseNODATA:
		return RespNODATA
	default:
		return RespError
	}
}

func (r *Response) hasFinalAnswer() bool {
	for _, rr := range r.Decoded.Answers {
		switch rr.(type) {
		case *miekgdns.CNAME, *miekgdns.DNAME, *miekgdns.RRSIG:
			// CNAME and DNAME are redirect records, not final answers.
			// RRSIG is a DNSSEC signature record — it covers the CNAME/DNAME
			// but is not itself the answer to the original question type.
			continue
		}
		return true
	}
	return false
}

func (r *Response) ChildReferrals() []*Referral {
	if r.Type != RespReferral {
		return nil
	}
	if r.Referral == nil {
		return nil
	}

	var nameservers []string
	for _, rr := range r.Decoded.Authority {
		if ns, ok := rr.(*miekgdns.NS); ok {
			if r.Referral.InBailiwick(ns.Ns) {
				nameservers = append(nameservers, ns.Ns)
			}
		}
	}

	if len(nameservers) == 0 {
		for _, rr := range r.Decoded.Authority {
			if ns, ok := rr.(*miekgdns.NS); ok {
				nameservers = append(nameservers, ns.Ns)
			}
		}
	}

	r.storeAuthority(nameservers)

	prob := r.childProb(len(nameservers))
	var children []*Referral
	for _, ns := range nameservers {
		child := NewReferral(
			r.Referral.Name,
			r.Referral.Qtype,
			ns,
			r.Referral.Depth+1,
			prob,
			r.Referral,
		)
		r.resolveGlue(child)
		children = append(children, child)
	}
	return children
}

func (r *Response) CNAMEFollowReferral() *Referral {
	if r.Type != RespCNAMEFollow || len(r.Decoded.CNAMEChain) == 0 {
		return nil
	}
	target := r.Decoded.CNAMEChain[len(r.Decoded.CNAMEChain)-1]
	follow := NewReferral(
		target,
		r.Referral.Qtype,
		r.Referral.Bailiwick,
		r.Referral.Depth+1,
		r.Referral.Prob,
		r.Referral,
	)
	if len(r.Referral.Addresses) > 0 {
		follow.Addresses = make([]net.IP, len(r.Referral.Addresses))
		copy(follow.Addresses, r.Referral.Addresses)
		follow.State = StateResolved
	}
	return follow
}

func (r *Response) storeAuthority(nameservers []string) {
	if r.Cache == nil {
		return
	}
	zone := r.Referral.Name
	r.Cache.StoreNS(zone, nameservers)
}

func (r *Response) resolveGlue(child *Referral) {
	if r.Cache == nil {
		return
	}
	nsName := child.Bailiwick
	for _, rr := range r.Decoded.Additional {
		switch v := rr.(type) {
		case *miekgdns.A:
			if normalize(v.Header().Name) == normalize(nsName) {
				child.Addresses = append(child.Addresses, v.A)
			}
		case *miekgdns.AAAA:
			if normalize(v.Header().Name) == normalize(nsName) {
				child.Addresses = append(child.Addresses, v.AAAA)
			}
		}
	}
	if child.HasAddresses() {
		child.State = StateResolved
	}
	r.Cache.StoreGlue(nsName, child.Addresses)
}

func (r *Response) IsTerminal() bool {
	switch r.Type {
	case RespAnswer, RespNODATA, RespNXDOMAIN, RespSERVFAIL, RespREFUSED, RespNOTIMPL, RespCNAMELoop, RespError, RespNSResolutionFailed:
		return true
	default:
		return false
	}
}

func (r *Response) childProb(n int) float64 {
	if n <= 0 {
		return 0
	}
	if r.Referral == nil {
		return 1.0 / float64(n)
	}
	return r.Referral.Prob / float64(n)
}