package traverse

import (
	"net"

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

// ResponseType classifies the outcome of querying a single Referral.
type ResponseType int

// Response type constants used to drive traversal logic.
const (
	RespReferral   ResponseType = iota // server returned an NS referral
	RespAnswer                         // server returned a final answer
	RespCNAMEFollow                    // answer contains a CNAME requiring further traversal
	RespNODATA                         // NOERROR with no matching records
	RespNXDOMAIN                       // name does not exist
	RespSERVFAIL                       // server failure
	RespREFUSED                        // query refused
	RespNOTIMPL                        // query type not implemented
	RespCNAMELoop                      // CNAME chain revisits a name already in the chain
	RespError                          // transport or decoding error
)

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"
	default:
		return "unknown"
	}
}

// Response is the result of querying a single Referral against a specific
// nameserver.  It contains the decoded DNS message, the classified ResponseType,
// and any error message for display.
type Response struct {
	// Referral is the query this response corresponds to.
	Referral *Referral
	// Server is the nameserver IP that was queried.
	Server net.IP
	// Cache is the InfoCache used during processing (for glue resolution).
	Cache *InfoCache
	// Decoded holds the structured DNS response fields.
	Decoded *dns.DecodedResponse
	// Type is the high-level classification of this response.
	Type ResponseType
	// ErrorMessage is a human-readable description when Type is RespError or RespCNAMELoop.
	ErrorMessage string
}

// NewResponse creates a Response for the given Referral and server.
func NewResponse(ref *Referral, server net.IP, cache *InfoCache) *Response {
	return &Response{
		Referral: ref,
		Server:   server,
		Cache:    cache,
	}
}

// Process decodes msg, classifies it, and populates r.Type and r.Decoded.
// Synthesises CNAME records from DNAME mappings when no explicit CNAME is present.
// Returns r for chaining.
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
}

// ChildReferrals returns the set of child Referrals implied by a referral
// response.  It extracts nameservers from the authority section, resolves
// any glue from the additional section, and sets Prob proportionally.
// Returns nil when Type != RespReferral.
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
}

// CNAMEFollowReferral constructs a follow-up Referral targeting the last CNAME
// in the chain.  Returns nil when Type != RespCNAMEFollow.
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)
}

// IsTerminal reports whether this response ends a traversal branch (no further
// referrals or CNAME follows are expected).
func (r *Response) IsTerminal() bool {
	switch r.Type {
	case RespAnswer, RespNODATA, RespNXDOMAIN, RespSERVFAIL, RespREFUSED, RespNOTIMPL, RespCNAMELoop, RespError:
		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)
}