package traverse

import (
	"context"
	"math"
	"net"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	idns "gitea.hansenits.com.au/hits/ExploreDNS/internal/dns"
	"github.com/miekg/dns"
)

// --- mock exchange: the single query path used by production and tests ---

type mockKey struct {
	server string
	qname  string
	qtype  uint16
}

type mockExchange struct {
	mu        sync.Mutex
	responses map[mockKey]*dns.Msg
	errors    map[mockKey]error
	calls     map[mockKey]int
}

func newMockExchange() *mockExchange {
	return &mockExchange{
		responses: make(map[mockKey]*dns.Msg),
		errors:    make(map[mockKey]error),
		calls:     make(map[mockKey]int),
	}
}

func (m *mockExchange) on(server, qname string, qtype uint16, msg *dns.Msg) {
	m.responses[mockKey{server, dns.Fqdn(qname), qtype}] = msg
}

func (m *mockExchange) fn(_ context.Context, server string, msg *dns.Msg, _ bool) (*dns.Msg, error) {
	host := server
	if h, _, err := net.SplitHostPort(server); err == nil {
		host = h
	}
	q := msg.Question[0]
	key := mockKey{host, q.Name, q.Qtype}
	m.mu.Lock()
	m.calls[key]++
	resp, ok := m.responses[key]
	err := m.errors[key]
	m.mu.Unlock()
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, &net.OpError{Op: "read", Err: &net.DNSError{Err: "no mock response", Name: q.Name}}
	}
	out := resp.Copy()
	out.SetReply(msg)
	out.Answer = resp.Answer
	out.Ns = resp.Ns
	out.Extra = resp.Extra
	out.Rcode = resp.Rcode
	return out, nil
}

func answerMsg(rrs ...dns.RR) *dns.Msg {
	m := new(dns.Msg)
	m.Answer = rrs
	return m
}

func referralMsg(nsRRs []dns.RR, glue ...dns.RR) *dns.Msg {
	m := new(dns.Msg)
	m.Ns = nsRRs
	m.Extra = glue
	return m
}

func rcodeMsg(rcode int) *dns.Msg {
	m := new(dns.Msg)
	m.Rcode = rcode
	return m
}

func testConfig(fast bool) *TraverserConfig {
	return &TraverserConfig{
		MaxDepth:  DefaultMaxDepth,
		QueryType: dns.TypeA,
		RootAddrs: []net.IP{net.ParseIP("198.41.0.4")},
		Fast:      fast,
		QueryConfig: &idns.QueryConfig{
			Retries:    1,
			Timeout:    time.Second,
			RetryDelay: time.Millisecond,
		},
	}
}

func runTraversal(t *testing.T, cfg *TraverserConfig, m *mockExchange, qname string) (*Traverser, *Referral) {
	t.Helper()
	tr := NewTraverser(cfg)
	tr.SetExchange(m.fn)
	root, err := tr.Run(context.Background(), qname)
	if err != nil {
		t.Fatalf("Run(%q): %v", qname, err)
	}
	return tr, root
}

func statsSum(root *Referral) float64 {
	sum := 0.0
	for _, e := range root.Stats {
		sum += e.Prob
	}
	return sum
}

func assertSumsToOne(t *testing.T, root *Referral) {
	t.Helper()
	if sum := statsSum(root); math.Abs(sum-1.0) > 1e-9 {
		t.Errorf("aggregated leaf probabilities sum to %v, want 1.0", sum)
	}
}

func leavesByStatus(root *Referral, status Status) []*StatsEntry {
	var out []*StatsEntry
	for _, e := range root.StatsList() {
		if e.Response.Status == status {
			out = append(out, e)
		}
	}
	return out
}

// --- scenarios ---

// mockSimpleDelegation wires root → com → example.com with two glued NS that
// both answer.
func mockSimpleDelegation() *mockExchange {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	))
	m.on("192.5.6.30", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com"), nsRR("example.com", "ns2.example.com")},
		aRR("ns1.example.com", "1.1.1.1"),
		aRR("ns2.example.com", "2.2.2.2"),
	))
	m.on("1.1.1.1", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))
	m.on("2.2.2.2", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))
	return m
}

func TestReferralFanOut(t *testing.T) {
	m := mockSimpleDelegation()
	_, root := runTraversal(t, testConfig(false), m, "www.example.com")

	assertSumsToOne(t, root)
	answered := leavesByStatus(root, StatusAnswered)
	if len(answered) != 2 {
		t.Fatalf("expected 2 answered leaves (one per NS), got %d: %v", len(answered), root.StatsList())
	}
	for _, leaf := range answered {
		if math.Abs(leaf.Prob-0.5) > 1e-9 {
			t.Errorf("leaf %s prob = %v, want 0.5", leaf.Key, leaf.Prob)
		}
	}

	// RefID grammar: rootroot "", root child "1", gtld "1.1", NS "1.1.1"/"1.1.2".
	if root.RefID != "" {
		t.Errorf("rootroot refid = %q, want empty", root.RefID)
	}
	top := root.Children["rootroot"]
	if len(top) != 1 || top[0].RefID != "1" {
		t.Fatalf("top children = %v", top)
	}
	gtld := top[0].Children["198.41.0.4"]
	if len(gtld) != 1 || gtld[0].RefID != "1.1" {
		t.Fatalf("gtld children refids wrong: %v", gtld)
	}
	nsKids := gtld[0].Children["192.5.6.30"]
	if len(nsKids) != 2 || nsKids[0].RefID != "1.1.1" || nsKids[1].RefID != "1.1.2" {
		t.Fatalf("NS children refids wrong: %v", nsKids)
	}
	if nsKids[0].Depth() != 3 {
		t.Errorf("depth of 1.1.1 = %d, want 3", nsKids[0].Depth())
	}
}

func TestGluelessResolveSubtree(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	))
	m.on("192.5.6.30", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com"), nsRR("example.com", "ns.other.net")},
		aRR("ns1.example.com", "1.1.1.1"),
	))
	m.on("1.1.1.1", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))
	// resolve subtree for A ns.other.net starts back at the root hints
	m.on("198.41.0.4", "ns.other.net", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("net", "d.gtld.net")},
		aRR("d.gtld.net", "3.3.3.3"),
	))
	m.on("3.3.3.3", "ns.other.net", dns.TypeA, answerMsg(aRR("ns.other.net", "4.4.4.4")))
	m.on("4.4.4.4", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))

	var events []TraversalEvent
	cfg := testConfig(false)
	cfg.Hooks = &TraverserHooks{OnEvent: func(ev TraversalEvent) { events = append(events, ev) }}
	_, root := runTraversal(t, cfg, m, "www.example.com")

	assertSumsToOne(t, root)
	answered := leavesByStatus(root, StatusAnswered)
	if len(answered) != 2 {
		t.Fatalf("expected 2 answered leaves, got %v", root.StatsList())
	}
	var viaGlueless *StatsEntry
	for _, leaf := range answered {
		if leaf.Response.IP == "4.4.4.4" {
			viaGlueless = leaf
		}
	}
	if viaGlueless == nil {
		t.Fatal("no answered leaf via the glueless nameserver")
	}
	if math.Abs(viaGlueless.Prob-0.5) > 1e-9 {
		t.Errorf("glueless leaf prob = %v, want 0.5", viaGlueless.Prob)
	}
	if got := viaGlueless.Referral.ServerWeights["4.4.4.4"]; math.Abs(got-1.0) > 1e-9 {
		t.Errorf("resolved serverweight = %v, want 1.0", got)
	}

	// The resolve subtree inserts a .0 refid component and is flagged.
	sawResolve := false
	for _, ev := range events {
		if ev.RefID == "1.1.2.0.1" {
			sawResolve = true
			if !ev.IsResolve {
				t.Error("resolve subtree event not flagged IsResolve")
			}
		}
	}
	if !sawResolve {
		t.Errorf("no event for resolve refid 1.1.2.0.1; events: %v", refids(events))
	}
	// Depth ignores the zero components.
	if d := refidDepth("1.1.2.0.1.1"); d != 5 {
		t.Errorf("refidDepth(1.1.2.0.1.1) = %d, want 5", d)
	}
}

func TestNoGlueDeadEnd(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	))
	// In-bailiwick NS without glue: dead end.
	m.on("192.5.6.30", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com")},
	))
	_, root := runTraversal(t, testConfig(false), m, "www.example.com")

	assertSumsToOne(t, root)
	noglue := leavesByStatus(root, StatusNoGlue)
	if len(noglue) != 1 {
		t.Fatalf("expected 1 noglue leaf, got %v", root.StatsList())
	}
	leaf := noglue[0]
	if math.Abs(leaf.Prob-1.0) > 1e-9 {
		t.Errorf("noglue prob = %v, want 1.0", leaf.Prob)
	}
	if leaf.Response.IP != "192.5.6.30" {
		t.Errorf("noglue response IP = %q, want the referring parent IP", leaf.Response.IP)
	}
	if leaf.Referral.Server != "ns1.example.com" {
		t.Errorf("noglue referral server = %q", leaf.Referral.Server)
	}
	if leaf.Referral.Parent.Server != "a.gtld-servers.net" {
		t.Errorf("noglue parent server = %q", leaf.Referral.Parent.Server)
	}
	if !strings.HasPrefix(leaf.Key, "key:noglue:192.5.6.30:www.example.com:IN:A:ns1.example.com:") {
		t.Errorf("noglue stats key = %q", leaf.Key)
	}
}

func TestResolveLoopDeadEnd(t *testing.T) {
	m := newMockExchange()
	// x.net NS ns.y.net (no glue); y.net NS ns.x.net (no glue): resolving
	// either server needs the other, which is a loop.
	xReferral := referralMsg([]dns.RR{nsRR("x.net", "ns.y.net")})
	yReferral := referralMsg([]dns.RR{nsRR("y.net", "ns.x.net")})
	m.on("198.41.0.4", "www.x.net", dns.TypeA, xReferral)
	m.on("198.41.0.4", "ns.y.net", dns.TypeA, yReferral)
	m.on("198.41.0.4", "ns.x.net", dns.TypeA, xReferral)

	_, root := runTraversal(t, testConfig(false), m, "www.x.net")

	assertSumsToOne(t, root)
	loops := leavesByStatus(root, StatusLoop)
	if len(loops) != 1 {
		t.Fatalf("expected 1 loop leaf, got %v", root.StatsList())
	}
	if math.Abs(loops[0].Prob-1.0) > 1e-9 {
		t.Errorf("loop prob = %v, want 1.0", loops[0].Prob)
	}
	if loops[0].Referral.Status != RefStatusLoop {
		t.Errorf("loop referral status = %q", loops[0].Referral.Status)
	}
}

func TestCNAMERestart(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.a.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	))
	m.on("192.5.6.30", "www.a.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("a.com", "ns.a.com")},
		aRR("ns.a.com", "5.5.5.5"),
	))
	m.on("5.5.5.5", "www.a.com", dns.TypeA, answerMsg(cnameRRT("www.a.com", "www.b.net")))
	// restart resumes from the deepest cached zone — nothing cached for
	// b.net, so back to the root.
	m.on("198.41.0.4", "www.b.net", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("net", "e.gtld.net")},
		aRR("e.gtld.net", "6.6.6.6"),
	))
	m.on("6.6.6.6", "www.b.net", dns.TypeA, answerMsg(aRR("www.b.net", "7.7.7.7")))

	_, root := runTraversal(t, testConfig(false), m, "www.a.com")

	assertSumsToOne(t, root)
	answered := leavesByStatus(root, StatusAnswered)
	if len(answered) != 1 {
		t.Fatalf("expected 1 answered leaf, got %v", root.StatsList())
	}
	leaf := answered[0]
	if leaf.Response.Qname != "www.b.net" {
		t.Errorf("answered qname = %q, want restart target www.b.net", leaf.Response.Qname)
	}
	if leaf.Referral.Qname != "www.b.net" {
		t.Errorf("restart referral qname = %q", leaf.Referral.Qname)
	}
	// The restart chain keeps numbering below the restarting node.
	if leaf.Referral.RefID != "1.1.1.1.1" {
		t.Errorf("answered refid = %q, want 1.1.1.1.1", leaf.Referral.RefID)
	}
}

func TestCNAMERestartLoop(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.a.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	))
	m.on("192.5.6.30", "www.a.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("a.com", "ns.a.com")},
		aRR("ns.a.com", "5.5.5.5"),
	))
	m.on("5.5.5.5", "www.a.com", dns.TypeA, answerMsg(cnameRRT("www.a.com", "www.b.net")))
	// www.b.net points straight back at www.a.com: every chain target is
	// checked against the ancestor queries, so this is a CNAME loop.
	m.on("198.41.0.4", "www.b.net", dns.TypeA, answerMsg(cnameRRT("www.b.net", "www.a.com")))

	_, root := runTraversal(t, testConfig(false), m, "www.a.com")

	assertSumsToOne(t, root)
	loops := leavesByStatus(root, StatusCNAMELoop)
	if len(loops) != 1 {
		t.Fatalf("expected 1 cname_loop leaf, got %v", root.StatsList())
	}
	if math.Abs(loops[0].Prob-1.0) > 1e-9 {
		t.Errorf("cname_loop prob = %v, want 1.0", loops[0].Prob)
	}
}

func TestCNAMERestartLoopIntermediateChainTarget(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.a.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("a.com", "ns.a.com")},
		aRR("ns.a.com", "5.5.5.5"),
	))
	m.on("5.5.5.5", "www.a.com", dns.TypeA, answerMsg(cnameRRT("www.a.com", "www.b.net")))
	m.on("198.41.0.4", "www.b.net", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("b.net", "ns.b.net")},
		aRR("ns.b.net", "6.6.6.6"),
	))
	// A three-record chain whose INTERMEDIATE target www.a.com matches an
	// ancestor query; the endname other.org does not, so an endname-only
	// loop check would miss it.
	m.on("6.6.6.6", "www.b.net", dns.TypeA, answerMsg(
		cnameRRT("www.b.net", "c.b.net"),
		cnameRRT("c.b.net", "www.a.com"),
		cnameRRT("www.a.com", "other.org"),
	))

	_, root := runTraversal(t, testConfig(false), m, "www.a.com")

	assertSumsToOne(t, root)
	loops := leavesByStatus(root, StatusCNAMELoop)
	if len(loops) != 1 {
		t.Fatalf("expected 1 cname_loop leaf, got %v", root.StatsList())
	}
	leaf := loops[0]
	if math.Abs(leaf.Prob-1.0) > 1e-9 {
		t.Errorf("cname_loop prob = %v, want 1.0", leaf.Prob)
	}
	dq := leaf.Response.DQ
	if len(dq.ChainTargets) != 3 {
		t.Fatalf("ChainTargets = %v, want 3 entries", dq.ChainTargets)
	}
	if dq.ChainTargets[1] != "www.a.com" {
		t.Errorf("ChainTargets[1] = %q, want www.a.com", dq.ChainTargets[1])
	}
	if dq.Endname != "other.org" {
		t.Errorf("Endname = %q, want other.org", dq.Endname)
	}
}

func TestDepthLimitInjectsException(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.c.b.a", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("a", "ns.a")},
		aRR("ns.a", "1.1.1.1"),
	))
	m.on("1.1.1.1", "www.c.b.a", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("b.a", "ns.b.a")},
		aRR("ns.b.a", "2.2.2.2"),
	))
	// Node 1.1.1 sits at depth 3 == maxdepth: its query is never sent.
	cfg := testConfig(false)
	cfg.MaxDepth = 3
	_, root := runTraversal(t, cfg, m, "www.c.b.a")

	assertSumsToOne(t, root)
	exceptions := leavesByStatus(root, StatusException)
	if len(exceptions) != 1 {
		t.Fatalf("expected 1 exception leaf, got %v", root.StatsList())
	}
	leaf := exceptions[0]
	if leaf.Response.DQ.ExceptionMessage != "Maxdepth 3 exceeded" {
		t.Errorf("exception message = %q, want %q", leaf.Response.DQ.ExceptionMessage, "Maxdepth 3 exceeded")
	}
	if math.Abs(leaf.Prob-1.0) > 1e-9 {
		t.Errorf("exception prob = %v, want 1.0", leaf.Prob)
	}
	// No query should have reached depth 3.
	m.mu.Lock()
	defer m.mu.Unlock()
	for key := range m.calls {
		if key.server == "2.2.2.2" {
			t.Error("query sent beyond the depth limit")
		}
	}
}

func TestLameReferral(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	))
	m.on("192.5.6.30", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com")},
		aRR("ns1.example.com", "1.1.1.1"),
	))
	// ns1 refers back to the same zone: not strictly deeper, so lame.
	m.on("1.1.1.1", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns2.example.com")},
		aRR("ns2.example.com", "2.2.2.2"),
	))
	_, root := runTraversal(t, testConfig(false), m, "www.example.com")

	assertSumsToOne(t, root)
	lame := leavesByStatus(root, StatusReferralLame)
	if len(lame) != 1 {
		t.Fatalf("expected 1 referral_lame leaf, got %v", root.StatsList())
	}
	leaf := lame[0]
	if math.Abs(leaf.Prob-1.0) > 1e-9 {
		t.Errorf("lame prob = %v, want 1.0", leaf.Prob)
	}
	if !strings.HasSuffix(leaf.Key, ":192.5.6.30") {
		t.Errorf("lame stats key should end with the parent IP, got %q", leaf.Key)
	}
	if leaf.Referral.ParentIP != "192.5.6.30" {
		t.Errorf("lame referral parent ip = %q", leaf.Referral.ParentIP)
	}
}

func TestChildsetDigitWhenMultipleIPsProduceChildren(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
		aRR("a.gtld-servers.net", "192.5.6.31"),
	))
	exampleReferral := referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com")},
		aRR("ns1.example.com", "1.1.1.1"),
	)
	m.on("192.5.6.30", "www.example.com", dns.TypeA, exampleReferral)
	m.on("192.5.6.31", "www.example.com", dns.TypeA, exampleReferral)
	m.on("1.1.1.1", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))

	var setEvents []TraversalEvent
	cfg := testConfig(false)
	cfg.Hooks = &TraverserHooks{OnEvent: func(ev TraversalEvent) {
		if ev.Stage == StageNewReferralSet {
			setEvents = append(setEvents, ev)
		}
	}}
	_, root := runTraversal(t, cfg, m, "www.example.com")

	assertSumsToOne(t, root)
	gtld := root.Children["rootroot"][0].Children["198.41.0.4"][0]
	set1 := gtld.Children["192.5.6.30"]
	set2 := gtld.Children["192.5.6.31"]
	if len(set1) != 1 || set1[0].RefID != "1.1.1.1" {
		t.Errorf("first childset refid = %v, want 1.1.1.1", refidsOf(set1))
	}
	if len(set2) != 1 || set2[0].RefID != "1.1.2.1" {
		t.Errorf("second childset refid = %v, want 1.1.2.1", refidsOf(set2))
	}
	if len(setEvents) != 2 {
		t.Errorf("expected 2 new_referral_set events, got %d", len(setEvents))
	}
	// Identical answers from both paths merge into one leaf with prob 1.0.
	answered := leavesByStatus(root, StatusAnswered)
	if len(answered) != 1 || math.Abs(answered[0].Prob-1.0) > 1e-9 {
		t.Errorf("answered leaves = %v", root.StatsList())
	}
}

func TestFastModeReuse(t *testing.T) {
	cfg := testConfig(true)
	cfg.RootAddrs = []net.IP{net.ParseIP("198.41.0.4"), net.ParseIP("199.9.14.201")}

	m := newMockExchange()
	comReferral := referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	)
	m.on("198.41.0.4", "www.example.com", dns.TypeA, comReferral)
	m.on("199.9.14.201", "www.example.com", dns.TypeA, comReferral)
	m.on("192.5.6.30", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))

	var events []TraversalEvent
	cfg.Hooks = &TraverserHooks{OnEvent: func(ev TraversalEvent) { events = append(events, ev) }}
	_, root := runTraversal(t, cfg, m, "www.example.com")

	assertSumsToOne(t, root)
	var fast *TraversalEvent
	for i := range events {
		if events[i].Stage == StageAnswerFast {
			fast = &events[i]
		}
	}
	if fast == nil {
		t.Fatal("expected a StageAnswerFast event in fast mode")
	}
	if fast.RefID != "2.1" || fast.CompletedEarlier != "1.1" {
		t.Errorf("fast event refid=%q completedEarlier=%q, want 2.1 / 1.1", fast.RefID, fast.CompletedEarlier)
	}
	if fast.Referral.ReplacedBy == nil || fast.Referral.ReplacedBy.RefID != "1.1" {
		t.Error("fast-replaced referral should point at its replacement")
	}
	// The second branch's child is the first branch's completed node.
	second := root.Children["rootroot"][1]
	if got := second.Children["199.9.14.201"][0].RefID; got != "1.1" {
		t.Errorf("replaced child refid = %q, want 1.1", got)
	}
	answered := leavesByStatus(root, StatusAnswered)
	if len(answered) != 1 || math.Abs(answered[0].Prob-1.0) > 1e-9 {
		t.Errorf("answered leaves = %v", root.StatsList())
	}
}

func TestNonFastModeReWalks(t *testing.T) {
	cfg := testConfig(false)
	cfg.RootAddrs = []net.IP{net.ParseIP("198.41.0.4"), net.ParseIP("199.9.14.201")}

	m := newMockExchange()
	comReferral := referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.net")},
		aRR("a.gtld-servers.net", "192.5.6.30"),
	)
	m.on("198.41.0.4", "www.example.com", dns.TypeA, comReferral)
	m.on("199.9.14.201", "www.example.com", dns.TypeA, comReferral)
	m.on("192.5.6.30", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "9.9.9.9")))

	var events []TraversalEvent
	cfg.Hooks = &TraverserHooks{OnEvent: func(ev TraversalEvent) { events = append(events, ev) }}
	_, root := runTraversal(t, cfg, m, "www.example.com")

	assertSumsToOne(t, root)
	for _, ev := range events {
		if ev.Stage == StageAnswerFast || ev.Stage == StageNewFast {
			t.Fatalf("unexpected fast-mode event %v in non-fast mode", ev.Stage)
		}
	}
	// Both branches keep their own child node.
	second := root.Children["rootroot"][1]
	if got := second.Children["199.9.14.201"][0].RefID; got != "2.1" {
		t.Errorf("non-fast child refid = %q, want 2.1", got)
	}
}

func TestAllRootsBranching(t *testing.T) {
	cfg := testConfig(false)
	cfg.RootAddrs = nil
	for i := 1; i <= 13; i++ {
		cfg.RootAddrs = append(cfg.RootAddrs, net.ParseIP("198.41.0."+strconv.Itoa(i)))
	}

	m := newMockExchange()
	for i := 1; i <= 13; i++ {
		m.on("198.41.0."+strconv.Itoa(i), "example.com", dns.TypeA, answerMsg(aRR("example.com", "9.9.9.9")))
	}
	_, root := runTraversal(t, cfg, m, "example.com")

	assertSumsToOne(t, root)
	top := root.Children["rootroot"]
	if len(top) != 13 {
		t.Fatalf("expected 13 top-level children, got %d", len(top))
	}
	for i, child := range top {
		if child.RefID != strconv.Itoa(i+1) {
			t.Errorf("child %d refid = %q, want %q", i, child.RefID, strconv.Itoa(i+1))
		}
	}
	answered := leavesByStatus(root, StatusAnswered)
	// Same answer from 13 different server IPs: 13 distinct leaves of 1/13.
	if len(answered) != 13 {
		t.Fatalf("expected 13 answered leaves, got %d", len(answered))
	}
	for _, leaf := range answered {
		if math.Abs(leaf.Prob-1.0/13) > 1e-9 {
			t.Errorf("leaf %s prob = %v, want %v", leaf.Key, leaf.Prob, 1.0/13)
		}
	}
}

func TestErrorAndNoDataStatuses(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com"), nsRR("example.com", "ns2.example.com")},
		aRR("ns1.example.com", "1.1.1.1"),
		aRR("ns2.example.com", "2.2.2.2"),
	))
	m.on("1.1.1.1", "www.example.com", dns.TypeA, rcodeMsg(dns.RcodeNameError))
	soa := &dns.SOA{
		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeSOA, Class: dns.ClassINET},
		Ns:  "ns1.example.com.", Mbox: "hostmaster.example.com.",
	}
	nodata := new(dns.Msg)
	nodata.Ns = []dns.RR{soa}
	m.on("2.2.2.2", "www.example.com", dns.TypeA, nodata)

	_, root := runTraversal(t, testConfig(false), m, "www.example.com")

	assertSumsToOne(t, root)
	errs := leavesByStatus(root, StatusError)
	if len(errs) != 1 || errs[0].Response.DQ.ErrorMessage != "No such domain (NXDOMAIN)" {
		t.Errorf("error leaves = %v", root.StatsList())
	}
	nodataLeaves := leavesByStatus(root, StatusNoData)
	if len(nodataLeaves) != 1 {
		t.Errorf("nodata leaves = %v", root.StatsList())
	}
}

func TestNetworkExceptionLeaf(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com")},
		aRR("ns1.example.com", "1.1.1.1"),
	))
	// no mock for 1.1.1.1 → network error → exception status
	_, root := runTraversal(t, testConfig(false), m, "www.example.com")

	assertSumsToOne(t, root)
	exceptions := leavesByStatus(root, StatusException)
	if len(exceptions) != 1 {
		t.Fatalf("expected exception leaf, got %v", root.StatsList())
	}
	if math.Abs(exceptions[0].Prob-1.0) > 1e-9 {
		t.Errorf("exception prob = %v", exceptions[0].Prob)
	}
}

func TestPacketCacheSingleWireQuery(t *testing.T) {
	m := mockSimpleDelegation()
	// Both NS answer; querying the same tuple twice must hit the cache.
	tr, _ := runTraversal(t, testConfig(false), m, "www.example.com")

	m.mu.Lock()
	defer m.mu.Unlock()
	for key, n := range m.calls {
		if n != 1 {
			t.Errorf("query %v sent %d times, want 1", key, n)
		}
	}
	if tr.client == nil {
		t.Fatal("traverser has no client")
	}
}

func TestServersEncountered(t *testing.T) {
	m := mockSimpleDelegation()
	tr, _ := runTraversal(t, testConfig(false), m, "www.example.com")
	seen := tr.ServersEncountered()
	if len(seen["ns1.example.com"]) != 1 || seen["ns1.example.com"][0] != "1.1.1.1" {
		t.Errorf("seen ns1 = %v", seen["ns1.example.com"])
	}
	if _, ok := seen["a.gtld-servers.net"]; !ok {
		t.Errorf("gtld server missing from seen: %v", seen)
	}
	if _, ok := seen[""]; ok {
		t.Error("rootroot must not be recorded in servers encountered")
	}
}

func TestRunContextCancellation(t *testing.T) {
	m := mockSimpleDelegation()
	tr := NewTraverser(testConfig(false))
	tr.SetExchange(m.fn)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := tr.Run(ctx, "www.example.com"); err == nil {
		t.Fatal("expected cancellation error")
	}
}

func TestIDNQnameConvertsToPunycode(t *testing.T) {
	m := newMockExchange()
	m.on("198.41.0.4", "xn--bcher-kva.example", dns.TypeA, answerMsg(aRR("xn--bcher-kva.example", "9.9.9.9")))
	_, root := runTraversal(t, testConfig(false), m, "bücher.example")
	if root.Qname != "xn--bcher-kva.example" {
		t.Errorf("qname = %q, want punycode", root.Qname)
	}
	assertSumsToOne(t, root)
}

func refids(events []TraversalEvent) []string {
	out := make([]string, len(events))
	for i, ev := range events {
		out[i] = ev.Stage.String() + ":" + ev.RefID
	}
	return out
}

func refidsOf(refs []*Referral) []string {
	out := make([]string, len(refs))
	for i, r := range refs {
		out[i] = r.RefID
	}
	return out
}

func cnameRRT(name, target string) dns.RR {
	return &dns.CNAME{
		Hdr:    dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeCNAME, Class: dns.ClassINET},
		Target: dns.Fqdn(target),
	}
}