package dns

import (
	"fmt"

	"github.com/miekg/dns"
)

type ResponseClassification int

const (
	ResponseAnswer ResponseClassification = iota
	ResponseReferral
	ResponseNODATA
	ResponseNXDOMAIN
	ResponseSERVFAIL
	ResponseOther
)

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

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         []string
}

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),
	}

	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.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 IsTruncated(msg *dns.Msg) bool {
	return msg != nil && msg.Truncated
}

func RcodeName(msg *dns.Msg) string {
	if msg == nil {
		return "UNKNOWN"
	}
	return dns.RcodeToString[msg.Rcode]
}

func ExtractAnswers(msg *dns.Msg) []dns.RR {
	if msg == nil {
		return nil
	}
	return msg.Answer
}

func ExtractAuthority(msg *dns.Msg) []dns.RR {
	if msg == nil {
		return nil
	}
	return msg.Ns
}

func ExtractCNAMEChain(msg *dns.Msg) []string {
	if msg == nil {
		return nil
	}
	return extractCNAMEChain(msg)
}

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
}

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
}

func HasCNAMEChain(msg *dns.Msg) bool {
	if msg == nil {
		return false
	}
	return len(extractCNAMEChain(msg)) > 0
}

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(),
	)
}