package traverse

import (
	"fmt"
	"strings"

	miekgdns "github.com/miekg/dns"
)

// Status is the classification of one query outcome. The first eight values
// come from decoded_query.rb / response.rb; noglue and loop are synthesised by
// the resolve phase without sending a query.
type Status string

const (
	StatusAnswered     Status = "answered"
	StatusNoData       Status = "nodata"
	StatusReferral     Status = "referral"
	StatusRestart      Status = "restart"
	StatusReferralLame Status = "referral_lame"
	StatusError        Status = "error"
	StatusException    Status = "exception"
	StatusCNAMELoop    Status = "cname_loop"
	StatusNoGlue       Status = "noglue"
	StatusLoop         Status = "loop"
)

// DecodedQuery classifies one DNS response (or network failure) against the
// query that produced it, mirroring decoded_query.rb. Names are canonical
// (lowercase, no trailing dot); Bailiwick "" means root.
type DecodedQuery struct {
	Msg *miekgdns.Msg
	Err error

	Qname     string
	Qclass    uint16
	Qtype     uint16
	IP        string
	Bailiwick string

	Status  Status
	Endname string
	// ChainTargets lists every CNAME target the in-message chain went
	// through (including the final unfollowed target when the chain leaves
	// the bailiwick); used for cross-response restart loop detection.
	ChainTargets []string

	CacheableGood []miekgdns.RR
	CacheableBad  []miekgdns.RR

	AuthNS    []miekgdns.RR
	AuthSOA   []miekgdns.RR
	AuthOther []miekgdns.RR

	Answers        []miekgdns.RR
	AuthorityNames []string

	ErrorMessage     string
	ExceptionMessage string
	Warnings         []string
}

// NewDecodedQuery decodes and classifies a response. Pass err non-nil for a
// network-level failure (dnstraverse's "exception"); msg is ignored then.
func NewDecodedQuery(msg *miekgdns.Msg, err error, qname string, qclass, qtype uint16, ip, bailiwick string) *DecodedQuery {
	dq := &DecodedQuery{
		Msg:       msg,
		Err:       err,
		Qname:     canonicalName(qname),
		Qclass:    qclass,
		Qtype:     qtype,
		IP:        ip,
		Bailiwick: canonicalName(bailiwick),
	}
	dq.process()
	return dq
}

func (dq *DecodedQuery) WarningsAdd(warnings ...string) {
	dq.Warnings = append(dq.Warnings, warnings...)
}

// process implements the classification order of decoded_query.rb#process
// exactly (the 7 steps in the design doc).
func (dq *DecodedQuery) process() {
	if dq.Err == nil && dq.Msg == nil {
		dq.Err = fmt.Errorf("nil DNS response")
	}
	if dq.Err != nil {
		dq.Status = StatusException
		dq.ExceptionMessage = dq.Err.Error()
		return
	}
	dq.AuthNS, dq.AuthSOA, dq.AuthOther = msgAuthority(dq.Msg)
	dq.CacheableGood, dq.CacheableBad = msgCacheable(dq.Msg, dq.Bailiwick)
	endname, targets, ok := msgFollowCNAMEs(dq.Msg, dq.Qname, dq.Qtype, dq.Bailiwick)
	if !ok {
		dq.Status = StatusCNAMELoop
		return
	}
	dq.Endname = endname
	dq.ChainTargets = targets
	if dq.Msg.Rcode != miekgdns.RcodeSuccess {
		dq.Status = StatusError
		dq.ErrorMessage = rcodeErrorMessage(dq.Msg.Rcode)
		return
	}
	if answers := msgAnswers(dq.Msg, dq.Endname, dq.Qclass, dq.Qtype); len(answers) > 0 {
		dq.Answers = answers
		dq.Status = StatusAnswered
		return
	}
	if dq.Endname != dq.Qname {
		dq.Status = StatusRestart
		return
	}
	if len(dq.AuthSOA) > 0 || len(dq.AuthNS) == 0 {
		dq.Status = StatusNoData
		return
	}
	dq.Status = StatusReferral
	for _, rr := range dq.AuthNS {
		if ns, ok := rr.(*miekgdns.NS); ok {
			dq.AuthorityNames = append(dq.AuthorityNames, canonicalName(ns.Ns))
		}
	}
}

// rcodeErrorMessage renders the exact error strings of decoded_query.rb
// process_error ("Format error" deliberately fixes the Ruby "Formate" typo —
// documented deviation).
func rcodeErrorMessage(rcode int) string {
	switch rcode {
	case miekgdns.RcodeFormatError:
		return "Format error (FORMERR)"
	case miekgdns.RcodeServerFailure:
		return "Server failure (SERVFAIL)"
	case miekgdns.RcodeNameError:
		return "No such domain (NXDOMAIN)"
	case miekgdns.RcodeNotImplemented:
		return "Not implemented (NOTIMP)"
	case miekgdns.RcodeRefused:
		return "Refused"
	default:
		if s, ok := miekgdns.RcodeToString[rcode]; ok {
			return s
		}
		return fmt.Sprintf("RCODE%d", rcode)
	}
}

// insideBailiwick reports whether name is at or below the bailiwick zone:
// bailiwick "" (root), equal fold, or name ends with "."+bailiwick.
func insideBailiwick(name, bailiwick string) bool {
	bw := canonicalName(bailiwick)
	if bw == "" {
		return true
	}
	n := canonicalName(name)
	return n == bw || strings.HasSuffix(n, "."+bw)
}

// msgAnswers returns the answer-section records matching qname/qclass/qtype
// (message_utility.rb msg_answers?). qtype ANY matches every type.
func msgAnswers(msg *miekgdns.Msg, qname string, qclass, qtype uint16) []miekgdns.RR {
	name := canonicalName(qname)
	any := qtype == miekgdns.TypeANY
	var out []miekgdns.RR
	for _, rr := range msg.Answer {
		h := rr.Header()
		if canonicalName(h.Name) == name && h.Class == qclass && (any || h.Rrtype == qtype) {
			out = append(out, rr)
		}
	}
	return out
}

// msgAuthority partitions the authority section into IN NS, IN SOA and other
// records (message_utility.rb msg_authority).
func msgAuthority(msg *miekgdns.Msg) (ns, soa, other []miekgdns.RR) {
	for _, rr := range msg.Ns {
		h := rr.Header()
		switch {
		case h.Rrtype == miekgdns.TypeNS && h.Class == miekgdns.ClassINET:
			ns = append(ns, rr)
		case h.Rrtype == miekgdns.TypeSOA && h.Class == miekgdns.ClassINET:
			soa = append(soa, rr)
		default:
			other = append(other, rr)
		}
	}
	return ns, soa, other
}

// msgCacheable partitions ALL sections (answer, authority, additional — in
// that order) into in-bailiwick records worth caching and out-of-bailiwick
// records to discard. OPT pseudo-records are dropped entirely.
func msgCacheable(msg *miekgdns.Msg, bailiwick string) (good, bad []miekgdns.RR) {
	for _, section := range [][]miekgdns.RR{msg.Answer, msg.Ns, msg.Extra} {
		for _, rr := range section {
			if rr.Header().Rrtype == miekgdns.TypeOPT {
				continue
			}
			if insideBailiwick(rr.Header().Name, bailiwick) {
				good = append(good, rr)
			} else {
				bad = append(bad, rr)
			}
		}
	}
	return good, bad
}

// msgFollowCNAMEs follows a CNAME chain within one message and returns the
// final name plus every target passed through (message_utility.rb
// msg_follow_cnames). Following stops — the target is returned unfollowed —
// as soon as the CURRENT owner name is not strictly below the bailiwick
// (Ruby tests `name !~ /\.#{bailiwick}$/i`, so an owner exactly equal to the
// bailiwick also stops the chain). An in-message loop returns ok=false
// (cname_loop).
func msgFollowCNAMEs(msg *miekgdns.Msg, qname string, qtype uint16, bailiwick string) (endname string, targets []string, ok bool) {
	name := canonicalName(qname)
	bw := canonicalName(bailiwick)
	seen := make(map[string]bool)
	for {
		seen[name] = true
		if len(msgAnswers(msg, name, miekgdns.ClassINET, qtype)) > 0 {
			return name, targets, true
		}
		cnames := msgAnswers(msg, name, miekgdns.ClassINET, miekgdns.TypeCNAME)
		if len(cnames) == 0 {
			return name, targets, true
		}
		cname, isCNAME := cnames[0].(*miekgdns.CNAME)
		if !isCNAME {
			return name, targets, true
		}
		target := canonicalName(cname.Target)
		targets = append(targets, target)
		if bw != "" && !strings.HasSuffix(name, "."+bw) {
			return target, targets, true
		}
		name = target
		if seen[name] {
			return "", targets, false
		}
	}
}

// isLameReferral implements the response.rb lame rule: a referral is lame
// unless the current bailiwick is root ("") or the new zone is STRICTLY
// deeper than the current bailiwick (equal or sideways zones are lame).
func isLameReferral(bailiwick, newBailiwick string) bool {
	bw := canonicalName(bailiwick)
	if bw == "" {
		return false
	}
	return !strings.HasSuffix(canonicalName(newBailiwick), "."+bw)
}