package traverse

import (
	"context"
	"fmt"
	"net"
	"sort"
	"strings"

	"gitea.hansenits.com.au/hits/ExploreDNS/internal/dns"
	miekgdns "github.com/miekg/dns"
	"golang.org/x/net/idna"
)

// DefaultMaxDepth is the default maximum referral depth (non-zero refid
// components) before a "Maxdepth N exceeded" exception is injected.
const DefaultMaxDepth = 20

// ReferralStatus is the resolve-phase status of a Referral node itself
// (referral.rb @status), distinct from the per-IP response statuses.
type ReferralStatus string

const (
	RefStatusNormal ReferralStatus = "normal"
	RefStatusNoGlue ReferralStatus = "noglue"
	RefStatusLoop   ReferralStatus = "loop"
)

// StatsEntry is one aggregated leaf statistic: the probability mass that
// ended in Response's outcome at Referral (referral.rb @stats values).
type StatsEntry struct {
	Key      string
	Prob     float64
	Response *ServerResponse
	Referral *Referral
}

// Referral represents one referral to a specific server for qname/qclass/
// qtype (referral.rb). The synthetic top node ("rootroot") has Server == ""
// and is never displayed; its children are the root servers.
type Referral struct {
	RefID  string
	Parent *Referral

	Qname  string
	Qclass uint16
	Qtype  uint16
	// NSAType is the record type used to resolve nameserver addresses
	// (always A; the reference is IPv4-only for transport).
	NSAType uint16

	// Server is the NS hostname this referral queries ("" for rootroot).
	Server string
	// ServerIPs is nil when the server still needs resolving. After a
	// resolve it may also contain "key:..." pseudo entries carrying the
	// probability of failed resolutions.
	ServerIPs []string
	Bailiwick string
	ParentIP  string

	InfoCache *InfoCache
	Status    ReferralStatus

	// Responses holds the classified response for each real IP queried.
	Responses map[string]*ServerResponse
	// Children holds child referrals keyed by the parent IP that produced
	// them ("rootroot" for the synthetic top node).
	Children map[string][]*Referral
	// Resolves is the glue-resolution subtree (refid ".0." components).
	Resolves []*Referral

	ServerWeights map[string]float64
	Warnings      []string

	// Stats is the post-order aggregation of leaf outcomes below this node.
	Stats map[string]*StatsEntry
	// StatsResolve aggregates the outcomes of the resolve subtree.
	StatsResolve map[string]*StatsEntry

	// ReplacedBy points at the earlier completed referral that fast mode
	// substituted for this node.
	ReplacedBy *Referral

	// summaryStats memoises SummaryStats() (referral.rb summary_stats).
	summaryStats *SummaryStats

	client             *dns.Client
	maxdepth           int
	referralResolution bool
	processed          bool
	calculated         bool
}

// idnaLookup is the IDN lookup profile used to convert internationalised
// domain names (unicode labels) to their ACE/punycode equivalents.
var idnaLookup = idna.New(
	idna.MapForLookup(),
	idna.BidiRule(),
	idna.StrictDomainName(false),
)

// toASCII converts a domain name that may contain unicode labels to its
// punycode (ACE) representation. On conversion errors the original name is
// returned so the caller can still attempt a query.
func toASCII(name string) string {
	if name == "" || name == "." {
		return name
	}
	ascii, err := idnaLookup.ToASCII(name)
	if err != nil {
		return name
	}
	return ascii
}

// referralArgs are the per-child overrides for makeReferral; zero values
// inherit from the parent (referral.rb make_referral merge semantics).
type referralArgs struct {
	qname              string
	qtype              uint16
	server             string
	serverIPs          []string
	bailiwick          string
	infoCache          *InfoCache
	refid              string
	parentIP           string
	referralResolution bool
}

func (r *Referral) makeReferral(a referralArgs) *Referral {
	child := &Referral{
		RefID:              a.refid,
		Parent:             r,
		Qname:              r.Qname,
		Qclass:             r.Qclass,
		Qtype:              r.Qtype,
		NSAType:            r.NSAType,
		Server:             canonicalName(a.server),
		ServerIPs:          a.serverIPs,
		Bailiwick:          canonicalName(a.bailiwick),
		ParentIP:           a.parentIP,
		InfoCache:          r.InfoCache,
		Status:             RefStatusNormal,
		Responses:          make(map[string]*ServerResponse),
		Children:           make(map[string][]*Referral),
		ServerWeights:      make(map[string]float64),
		client:             r.client,
		maxdepth:           r.maxdepth,
		referralResolution: a.referralResolution || r.referralResolution,
	}
	if a.qname != "" {
		child.Qname = canonicalName(a.qname)
	}
	if a.qtype != 0 {
		child.Qtype = a.qtype
	}
	if a.infoCache != nil {
		child.InfoCache = a.infoCache
	}
	// serverweight = 1/len(ips) per IP when the addresses are known.
	if child.ServerIPs != nil {
		for _, ip := range child.ServerIPs {
			child.ServerWeights[ip] = 1.0 / float64(len(child.ServerIPs))
		}
	}
	return child
}

// IsRootRoot reports whether this is the synthetic top node representing an
// automatic referral to the root servers.
func (r *Referral) IsRootRoot() bool {
	return r.Server == ""
}

// IsResolve reports whether this node is part of a glue-resolution subtree.
func (r *Referral) IsResolve() bool {
	return r.referralResolution
}

// Resolved reports whether the server addresses are known (rootroot is
// always resolved).
func (r *Referral) Resolved() bool {
	return r.IsRootRoot() || r.ServerIPs != nil
}

// Depth counts the non-zero refid components; resolve subtrees (".0.") do
// not count against the depth limit.
func (r *Referral) Depth() int {
	return refidDepth(r.RefID)
}

func refidDepth(refid string) int {
	if refid == "" {
		return 0
	}
	n := 0
	for _, part := range strings.Split(refid, ".") {
		if part != "0" {
			n++
		}
	}
	return n
}

// IPsAsArray returns the real IP addresses known for this referral,
// excluding "key:" pseudo entries.
func (r *Referral) IPsAsArray() []string {
	var out []string
	for _, ip := range r.ServerIPs {
		if !strings.HasPrefix(ip, "key:") {
			out = append(out, ip)
		}
	}
	return out
}

// TxtIPsVerbose renders the per-IP weights, sorted, e.g.
// "50.0%=1.2.3.4,50.0%=noglue:1.2.3.4" (referral.rb txt_ips_verbose). It is
// part of the fast-mode memo key.
func (r *Referral) TxtIPsVerbose() string {
	if r.ServerIPs == nil {
		return ""
	}
	parts := make([]string, 0, len(r.ServerIPs))
	for _, ip := range r.ServerIPs {
		label := ip
		if rest, ok := strings.CutPrefix(ip, "key:"); ok {
			// keep the first two colon-separated fields, like Ruby's
			// /^key:([^:]+(:[^:]*)?)/ capture.
			fields := strings.SplitN(rest, ":", 3)
			if len(fields) > 2 {
				fields = fields[:2]
			}
			label = strings.Join(fields, ":")
		}
		parts = append(parts, fmt.Sprintf("%.1f%%=%s", 100*r.ServerWeights[ip], label))
	}
	sort.Strings(parts)
	return strings.Join(parts, ",")
}

// TxtIPs renders the addresses for progress display; failed-resolve pseudo
// entries render as their response description (referral.rb txt_ips).
func (r *Referral) TxtIPs() string {
	if r.ServerIPs == nil {
		return ""
	}
	parts := make([]string, 0, len(r.ServerIPs))
	for _, ip := range r.ServerIPs {
		if strings.HasPrefix(ip, "key:") {
			if e, ok := r.StatsResolve[ip]; ok && e.Response != nil {
				parts = append(parts, e.Response.String())
				continue
			}
		}
		parts = append(parts, ip)
	}
	sort.Strings(parts)
	return strings.Join(parts, ",")
}

func (r *Referral) String() string {
	return fmt.Sprintf("%s [%s/%s/%s] server=%s server_ips=%s bailiwick=%s",
		r.RefID, r.Qname, ClassToString(r.Qclass), TypeToString(r.Qtype),
		r.Server, r.TxtIPs(), r.Bailiwick)
}

// OverallStatus summarises the node's outcome for event consumers: a resolve
// dead end (noglue/loop), the shared status of every per-IP response, or
// "mixed" when the responses disagree ("" before anything was queried).
func (r *Referral) OverallStatus() Status {
	switch r.Status {
	case RefStatusNoGlue:
		return StatusNoGlue
	case RefStatusLoop:
		return StatusLoop
	}
	var s Status
	for _, resp := range r.Responses {
		if s == "" {
			s = resp.Status
		} else if s != resp.Status {
			return "mixed"
		}
	}
	return s
}

// StatsList returns the aggregated leaf statistics sorted by stats key.
func (r *Referral) StatsList() []*StatsEntry {
	out := make([]*StatsEntry, 0, len(r.Stats))
	for _, e := range r.Stats {
		out = append(out, e)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Key < out[j].Key })
	return out
}

// isNoGlue reports a dead end: the server is inside the current bailiwick,
// so its address should have come as glue, but none was provided and no
// deeper zone can name it (referral.rb noglue?).
func (r *Referral) isNoGlue() bool {
	return r.ServerIPs == nil && insideBailiwick(r.Server, r.Bailiwick)
}

// isLoop reports a resolve loop: an ancestor referral asks the same
// qname/qclass/qtype of the same still-unresolved server (referral.rb loop?),
// e.g. b NS c.d while d NS a.b.
func (r *Referral) isLoop() bool {
	if r.ServerIPs != nil {
		return false
	}
	for p := r.Parent; p != nil; p = p.Parent {
		if p.Qname == r.Qname && p.Qclass == r.Qclass && p.Qtype == r.Qtype &&
			p.Server == r.Server && p.ServerIPs == nil {
			return true
		}
	}
	return false
}

// chainHasQuery reports whether this referral or any ancestor already asks
// qname with the same qclass/qtype. Used to stop cross-response CNAME chains
// (restart loops) that Ruby only catches via the depth limit.
func (r *Referral) chainHasQuery(qname string) bool {
	name := canonicalName(qname)
	for p := r; p != nil; p = p.Parent {
		if p.Qname == name && p.Qclass == r.Qclass && p.Qtype == r.Qtype {
			return true
		}
	}
	return false
}

// resolve turns an address-less referral into either a dead end (noglue/
// loop) or a resolve subtree querying A <server> from this branch's cache
// (referral.rb resolve). It returns the referrals to process.
func (r *Referral) resolve() ([]*Referral, error) {
	if r.isNoGlue() {
		r.Status = RefStatusNoGlue
		return nil, nil
	}
	if r.isLoop() {
		r.Status = RefStatusLoop
		return nil, nil
	}
	starters, newbailiwick, err := r.InfoCache.GetStartServers(r.Server)
	if err != nil {
		return nil, err
	}
	for i, st := range starters {
		child := r.makeReferral(referralArgs{
			qname:              r.Server,
			qtype:              r.NSAType,
			server:             st.Name,
			serverIPs:          st.IPs,
			bailiwick:          newbailiwick,
			refid:              fmt.Sprintf("%s.0.%d", r.RefID, i+1),
			referralResolution: true,
		})
		r.Resolves = append(r.Resolves, child)
	}
	return r.Resolves, nil
}

// resolveCalculate folds the resolve subtree's statistics into per-IP server
// weights (referral.rb resolve_calculate): each answered leaf distributes
// its probability evenly across the A records it returned; every other leaf
// keeps its probability under its "key:" stats key so failures surface in
// the results.
func (r *Referral) resolveCalculate() {
	r.StatsResolve = make(map[string]*StatsEntry)
	switch r.Status {
	case RefStatusNoGlue:
		resp := NewNoGlueResponse(r.Qname, r.Qclass, r.Qtype, r.ParentIP, r.Server, r.Bailiwick)
		key := resp.StatsKey()
		r.StatsResolve[key] = &StatsEntry{Key: key, Prob: 1.0, Response: resp, Referral: r}
	case RefStatusLoop:
		resp := NewLoopResponse(r.Qname, r.Qclass, r.Qtype, r.ParentIP, r.Server, r.Bailiwick)
		key := resp.StatsKey()
		r.StatsResolve[key] = &StatsEntry{Key: key, Prob: 1.0, Response: resp, Referral: r}
	default:
		statsCalculateChildren(r.StatsResolve, r.Resolves, 1.0)
	}

	r.ServerWeights = make(map[string]float64)
	r.ServerIPs = []string{}
	keys := make([]string, 0, len(r.StatsResolve))
	for key := range r.StatsResolve {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	addWeight := func(ip string, prob float64) {
		if _, ok := r.ServerWeights[ip]; !ok {
			r.ServerIPs = append(r.ServerIPs, ip)
		}
		r.ServerWeights[ip] += prob
	}
	for _, key := range keys {
		data := r.StatsResolve[key]
		if data.Response.Status == StatusAnswered {
			var addrs []string
			for _, rr := range data.Response.DQ.Answers {
				if a, ok := rr.(*miekgdns.A); ok {
					addrs = append(addrs, a.A.String())
				}
			}
			for _, addr := range addrs {
				addWeight(addr, data.Prob/float64(len(addrs)))
			}
			if len(addrs) == 0 {
				// answered but no A records (e.g. AAAA-only): carry the
				// probability as a failure key so mass is not lost.
				addWeight(key, data.Prob)
			}
		} else {
			addWeight(key, data.Prob)
		}
	}
}

// statsCalculateChildren merges the children's statistics into stats with an
// equal split of weight among them (referral.rb stats_calculate_children).
func statsCalculateChildren(stats map[string]*StatsEntry, children []*Referral, weight float64) {
	if len(children) == 0 {
		return
	}
	percent := (1.0 / float64(len(children))) * weight
	for _, child := range children {
		for key, data := range child.Stats {
			if e, ok := stats[key]; ok {
				e.Prob += data.Prob * percent
			} else {
				stats[key] = &StatsEntry{
					Key:      key,
					Prob:     data.Prob * percent,
					Response: data.Response,
					Referral: data.Referral,
				}
			}
		}
	}
}

// answerCalculate computes this node's aggregated statistics from its
// children, responses and resolve failures (referral.rb answer_calculate).
// Unlike the Ruby source, duplicate stats keys across a referral's IPs merge
// by summing probability (Ruby computes the sum then discards it — a source
// bug that breaks the probabilities-sum-to-1 invariant).
func (r *Referral) answerCalculate() {
	r.Stats = make(map[string]*StatsEntry)
	if r.IsRootRoot() {
		statsCalculateChildren(r.Stats, r.Children["rootroot"], 1.0)
		r.calculated = true
		return
	}
	for _, ip := range r.ServerIPs {
		serverweight := r.ServerWeights[ip]
		if strings.HasPrefix(ip, "key:") {
			// resolve failed for some reason - copy the resolve statistics
			src := r.StatsResolve[ip]
			if e, ok := r.Stats[ip]; ok {
				e.Prob += src.Prob
			} else {
				r.Stats[ip] = &StatsEntry{Key: ip, Prob: src.Prob, Response: src.Response, Referral: src.Referral}
			}
			continue
		}
		if children := r.Children[ip]; len(children) > 0 {
			statsCalculateChildren(r.Stats, children, serverweight)
			continue
		}
		resp := r.Responses[ip]
		if resp == nil {
			continue
		}
		key := resp.StatsKey()
		if e, ok := r.Stats[key]; ok {
			e.Prob += serverweight
		} else {
			r.Stats[key] = &StatsEntry{Key: key, Prob: serverweight, Response: resp, Referral: r}
		}
	}
	r.calculated = true
}

// process queries every real IP of this referral through the packet cache,
// classifies each response, and creates one child per NS name (including
// glueless ones) for referral/restart statuses (referral.rb process/
// process_normal). It returns one set of children per IP that produced any.
func (r *Referral) process(ctx context.Context) ([][]*Referral, error) {
	r.processed = true
	if r.IsRootRoot() {
		children, err := r.processAddRoots()
		if err != nil {
			return nil, err
		}
		return [][]*Referral{children}, nil
	}

	// Phase one: query and classify, counting childsets so refids can grow
	// an extra childset digit when more than one IP produces children.
	childsets := 0
	var order []string
	for _, ip := range r.ServerIPs {
		if strings.HasPrefix(ip, "key:") {
			continue
		}
		var dq *DecodedQuery
		if r.Depth() >= r.maxdepth {
			err := fmt.Errorf("Maxdepth %d exceeded", r.maxdepth)
			dq = NewDecodedQuery(nil, err, r.Qname, r.Qclass, r.Qtype, ip, r.Bailiwick)
		} else {
			msg, warnings, err := r.client.Query(ctx, net.ParseIP(ip), r.Qname, r.Qtype)
			dq = NewDecodedQuery(msg, err, r.Qname, r.Qclass, r.Qtype, ip, r.Bailiwick)
			dq.WarningsAdd(warnings...)
		}
		resp, err := NewServerResponse(dq, r.Server, r.ParentIP, r.InfoCache)
		if err != nil {
			return nil, err
		}
		if resp.Status == StatusRestart {
			// Cross-response CNAME loop: any target in the chain that we
			// (or an ancestor) are already querying is a dead end.
			for _, target := range dq.ChainTargets {
				if r.chainHasQuery(target) {
					resp.Status = StatusCNAMELoop
					break
				}
			}
		}
		r.Warnings = append(r.Warnings, dq.Warnings...)
		r.Responses[ip] = resp
		order = append(order, ip)
		if resp.Status == StatusRestart || resp.Status == StatusReferral {
			childsets++
		}
	}

	// Phase two: create the children.
	childset := 0
	var sets [][]*Referral
	for _, ip := range order {
		resp := r.Responses[ip]
		if resp.Status != StatusRestart && resp.Status != StatusReferral {
			continue
		}
		childset++
		refid := r.RefID
		if childsets > 1 {
			refid = fmt.Sprintf("%s.%d", r.RefID, childset)
		}
		children := r.makeReferrals(resp, refid, ip)
		r.Children[ip] = children
		sets = append(sets, children)
	}
	return sets, nil
}

// processAddRoots creates one child per root server with equal weight
// (referral.rb process_add_roots); the roots come from the info cache hints.
func (r *Referral) processAddRoots() ([]*Referral, error) {
	starters, _, err := r.InfoCache.GetStartServers("")
	if err != nil {
		return nil, err
	}
	dot := ""
	if r.RefID != "" {
		dot = "."
	}
	var children []*Referral
	for i, root := range starters {
		child := r.makeReferral(referralArgs{
			server:    root.Name,
			serverIPs: root.IPs,
			refid:     fmt.Sprintf("%s%s%d", r.RefID, dot, i+1),
		})
		children = append(children, child)
	}
	r.Children["rootroot"] = children
	return children, nil
}

// makeReferrals creates one child per start server for a referral/restart
// response (referral.rb make_referrals): qname moves to the response's
// endname (the CNAME target on restart), the bailiwick and cache come from
// the response.
func (r *Referral) makeReferrals(resp *ServerResponse, refid, parentIP string) []*Referral {
	var children []*Referral
	for i, st := range resp.Starters {
		children = append(children, r.makeReferral(referralArgs{
			qname:     resp.DQ.Endname,
			server:    st.Name,
			serverIPs: st.IPs,
			bailiwick: resp.StartersBailiwick,
			infoCache: resp.Cache,
			refid:     fmt.Sprintf("%s.%d", refid, i+1),
			parentIP:  parentIP,
		}))
	}
	return children
}

// replaceChild swaps before for after in the children/resolves lists (fast
// mode substitution); before keeps a pointer to its replacement.
func (r *Referral) replaceChild(before, after *Referral) {
	before.ReplacedBy = after
	for ip := range r.Children {
		for i, c := range r.Children[ip] {
			if c == before {
				r.Children[ip][i] = after
			}
		}
	}
	for i, c := range r.Resolves {
		if c == before {
			r.Resolves[i] = after
		}
	}
}