package dns

import (
	"fmt"
	"strings"

	"github.com/miekg/dns"
)

// ResponseClassification categorises a DNS response at a high level,
// independent of the raw RCODE.
type ResponseClassification int

// Classification constants, in order from most to least specific.
const (
	ResponseAnswer    ResponseClassification = iota // answer section contains records
	ResponseReferral                                // non-authoritative NS referral
	ResponseNODATA                                  // NOERROR with empty answer section
	ResponseNXDOMAIN                                // name does not exist (RCODE 3)
	ResponseSERVFAIL                                // server failure (RCODE 2)
	ResponseREFUSED                                 // query refused (RCODE 5)
	ResponseNOTIMPL                                 // not implemented (RCODE 4)
	ResponseOther                                   // any other RCODE
)

func (rc ResponseClassification) String() string {
	switch rc {
	case ResponseAnswer:
		return "answer"
	case ResponseReferral:
		return "referral"
	case ResponseNODATA:
		return "nodata"
	case ResponseNXDOMAIN:
		return "nxdomain"
	case ResponseSERVFAIL:
		return "servfail"
	case ResponseREFUSED:
		return "refused"
	case ResponseNOTIMPL:
		return "notimp"
	default:
		return "other"
	}
}

// DecodedResponse holds the extracted, structured fields from a raw *dns.Msg.
type DecodedResponse struct {
	Rcode              int
	RcodeName          string
	Truncated          bool
	RecursionAvailable bool
	Authoritative      bool
	Classification     ResponseClassification
	Answers            []dns.RR
	Authority          []dns.RR
	Additional         []dns.RR
	// CNAMEChain contains the ordered CNAME targets from the answer section.
	CNAMEChain []string
	// DNAMEMappings contains any DNAME records for redirect synthesis.
	DNAMEMappings []DNAMEMapping
}

// DNAMEMapping holds a DNAME record's owner and target for redirect synthesis.
type DNAMEMapping struct {
	Owner  string // e.g., "example.com."
	Target string // e.g., "example.net."
}

// DecodeResponse converts a raw *dns.Msg into a DecodedResponse.
// Returns nil when msg is nil.
func DecodeResponse(msg *dns.Msg) *DecodedResponse {
	if msg == nil {
		return nil
	}

	d := &DecodedResponse{
		Rcode:              msg.Rcode,
		RcodeName:          dns.RcodeToString[msg.Rcode],
		Truncated:          msg.Truncated,
		RecursionAvailable: msg.RecursionAvailable,
		Authoritative:      msg.Authoritative,
		Answers:            msg.Answer,
		Authority:          msg.Ns,
		Additional:         msg.Extra,
		CNAMEChain:         extractCNAMEChain(msg),
		DNAMEMappings:      extractDNAMEMappings(msg),
	}

	d.Classification = classify(msg)

	return d
}

func classify(msg *dns.Msg) ResponseClassification {
	switch msg.Rcode {
	case dns.RcodeNameError:
		return ResponseNXDOMAIN
	case dns.RcodeServerFailure:
		return ResponseSERVFAIL
	case dns.RcodeRefused:
		return ResponseREFUSED
	case dns.RcodeNotImplemented:
		return ResponseNOTIMPL
	case dns.RcodeSuccess:
		return classifySuccess(msg)
	default:
		return ResponseOther
	}
}

func classifySuccess(msg *dns.Msg) ResponseClassification {
	hasAnswers := len(msg.Answer) > 0

	if hasAnswers {
		return ResponseAnswer
	}

	hasNS := hasNSRecords(msg.Ns)
	if hasNS && !msg.Authoritative {
		return ResponseReferral
	}

	if hasNS {
		return ResponseNODATA
	}

	return ResponseNODATA
}

func hasNSRecords(rrs []dns.RR) bool {
	for _, rr := range rrs {
		if _, ok := rr.(*dns.NS); ok {
			return true
		}
	}
	return false
}

func extractCNAMEChain(msg *dns.Msg) []string {
	var chain []string
	seen := make(map[string]bool)
	for _, rr := range msg.Answer {
		if cname, ok := rr.(*dns.CNAME); ok {
			target := cname.Target
			if !seen[target] {
				seen[target] = true
				chain = append(chain, target)
			}
		}
	}
	return chain
}

func extractDNAMEMappings(msg *dns.Msg) []DNAMEMapping {
	var mappings []DNAMEMapping
	for _, rr := range msg.Answer {
		if dname, ok := rr.(*dns.DNAME); ok {
			mappings = append(mappings, DNAMEMapping{
				Owner:  dns.Fqdn(dname.Hdr.Name),
				Target: dns.Fqdn(dname.Target),
			})
		}
	}
	return mappings
}

// SynthesizeCNAMEFromDNAME computes the CNAME target for queryName given a DNAME mapping.
// Returns empty string if queryName is not a strict subdomain of dnameOwner.
func SynthesizeCNAMEFromDNAME(queryName, dnameOwner, dnameTarget string) string {
	q := strings.ToLower(dns.Fqdn(queryName))
	owner := strings.ToLower(dns.Fqdn(dnameOwner))
	target := strings.ToLower(dns.Fqdn(dnameTarget))

	if !dns.IsSubDomain(owner, q) || q == owner {
		return ""
	}
	prefix := strings.TrimSuffix(q, owner)
	return prefix + target
}

// IsTruncated reports whether msg has the TC (truncated) bit set.
func IsTruncated(msg *dns.Msg) bool {
	return msg != nil && msg.Truncated
}

// RcodeName returns the string representation of the RCODE in msg (e.g. "NOERROR").
// Returns "UNKNOWN" when msg is nil.
func RcodeName(msg *dns.Msg) string {
	if msg == nil {
		return "UNKNOWN"
	}
	return dns.RcodeToString[msg.Rcode]
}

// ExtractAnswers returns the answer section of msg, or nil when msg is nil.
func ExtractAnswers(msg *dns.Msg) []dns.RR {
	if msg == nil {
		return nil
	}
	return msg.Answer
}

// ExtractAuthority returns the authority section of msg, or nil when msg is nil.
func ExtractAuthority(msg *dns.Msg) []dns.RR {
	if msg == nil {
		return nil
	}
	return msg.Ns
}

// ExtractCNAMEChain returns the ordered list of CNAME targets from the answer section.
// Returns nil when msg is nil.
func ExtractCNAMEChain(msg *dns.Msg) []string {
	if msg == nil {
		return nil
	}
	return extractCNAMEChain(msg)
}

// IsReferral reports whether msg is a non-authoritative NS referral
// (NOERROR, empty answer section, NS records in authority, AA=false).
func IsReferral(msg *dns.Msg) bool {
	if msg == nil || msg.Rcode != dns.RcodeSuccess || len(msg.Answer) > 0 {
		return false
	}
	return hasNSRecords(msg.Ns) && !msg.Authoritative
}

// IsNODATA reports whether msg is a NODATA response
// (NOERROR with an empty answer section and no referral).
func IsNODATA(msg *dns.Msg) bool {
	if msg == nil || msg.Rcode != dns.RcodeSuccess {
		return false
	}
	if len(msg.Answer) > 0 {
		return false
	}
	if IsReferral(msg) {
		return false
	}
	return true
}

// HasCNAMEChain reports whether msg contains at least one CNAME record in its
// answer section.
func HasCNAMEChain(msg *dns.Msg) bool {
	if msg == nil {
		return false
	}
	return len(extractCNAMEChain(msg)) > 0
}

// FormatRecord returns a human-readable representation of a DNS resource record
// in the form: NAME TTL CLASS TYPE RDATA.
// Returns an empty string when rr is nil.
func FormatRecord(rr dns.RR) string {
	if rr == nil {
		return ""
	}
	header := rr.Header()
	return fmt.Sprintf("%s %d %s %s %s",
		header.Name,
		header.Ttl,
		dns.ClassToString[header.Class],
		QNameType(header.Rrtype),
		rr.String(),
	)
}