package traverse

import (
	"math"
	"net"
	"strconv"
	"strings"
	"testing"

	"github.com/miekg/dns"
)

// mockCaptureTopology reproduces the delegation behind
// docs/captures/dnstraverse-ruby-www.example.com-A.txt: one root, thirteen
// com gTLD servers, example.com served by two NS with three IPs each, every
// endpoint answering the same two A records (two endpoints return them in
// the opposite order, as in the capture).
func mockCaptureTopology() *mockExchange {
	m := newMockExchange()

	gtlds := []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m"}
	var comNS, comGlue []dns.RR
	gtldIPs := make([]string, len(gtlds))
	for i, l := range gtlds {
		ip := "192.0.2." + strconv.Itoa(i+1)
		gtldIPs[i] = ip
		comNS = append(comNS, nsRR("com", l+".gtld-servers.net"))
		comGlue = append(comGlue, aRR(l+".gtld-servers.net", ip))
	}
	m.on("202.12.27.33", "www.example.com", dns.TypeA, referralMsg(comNS, comGlue...))

	heraIPs := []string{"108.162.192.162", "172.64.32.162", "173.245.58.162"}
	elliottIPs := []string{"108.162.195.228", "162.159.44.228", "172.64.35.228"}
	exampleReferral := referralMsg(
		[]dns.RR{
			nsRR("example.com", "hera.ns.cloudflare.com"),
			nsRR("example.com", "elliott.ns.cloudflare.com"),
		},
		aRR("hera.ns.cloudflare.com", heraIPs[0]),
		aRR("hera.ns.cloudflare.com", heraIPs[1]),
		aRR("hera.ns.cloudflare.com", heraIPs[2]),
		aRR("elliott.ns.cloudflare.com", elliottIPs[0]),
		aRR("elliott.ns.cloudflare.com", elliottIPs[1]),
		aRR("elliott.ns.cloudflare.com", elliottIPs[2]),
	)
	for _, ip := range gtldIPs {
		m.on(ip, "www.example.com", dns.TypeA, exampleReferral)
	}

	forward := answerMsg(
		aRR("www.example.com", "104.20.23.154"),
		aRR("www.example.com", "172.66.147.243"),
	)
	reversed := answerMsg(
		aRR("www.example.com", "172.66.147.243"),
		aRR("www.example.com", "104.20.23.154"),
	)
	for _, ip := range []string{heraIPs[0], elliottIPs[0], elliottIPs[1], elliottIPs[2]} {
		m.on(ip, "www.example.com", dns.TypeA, forward)
	}
	for _, ip := range []string{heraIPs[1], heraIPs[2]} {
		m.on(ip, "www.example.com", dns.TypeA, reversed)
	}
	return m
}

// TestCapturePerEndpointFractions asserts the 16.7%-per-endpoint result of
// the www.example.com capture: 13 gTLD paths collapse (fast mode) into six
// endpoint leaves of 1/6 each, and the summary merges the differently
// ordered RRsets into a single 100% answered line.
func TestCapturePerEndpointFractions(t *testing.T) {
	m := mockCaptureTopology()
	cfg := testConfig(true)
	cfg.RootAddrs = []net.IP{net.ParseIP("202.12.27.33")}
	_, root := runTraversal(t, cfg, m, "www.example.com")

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

	stats := root.SummaryStats()
	if stats == nil {
		t.Fatal("expected summary stats after calculation")
	}
	if prob := stats.ByStatus[StatusAnswered]; math.Abs(prob-1.0) > 1e-9 {
		t.Errorf("answered summary prob = %v, want 1.0", prob)
	}
	// Both RR orders share the same sorted-rdata key: one summary line.
	if len(stats.Answers) != 1 {
		t.Fatalf("expected 1 distinct answered RRset, got %d: %v", len(stats.Answers), stats.Answers)
	}
	if math.Abs(stats.Answers[0].Prob-1.0) > 1e-9 {
		t.Errorf("answer group prob = %v, want 1.0", stats.Answers[0].Prob)
	}
	if !strings.Contains(stats.Answers[0].Key, "@@@") {
		t.Errorf("answer key should join rdata with @@@, got %q", stats.Answers[0].Key)
	}
	if len(stats.Answers[0].RRs) != 2 {
		t.Errorf("answer RRs = %v", stats.Answers[0].RRs)
	}
}

// TestDistinctRRsetsSeparateGroups asserts the converse of the capture test:
// two endpoints answering DIFFERENT content produce two separate summary
// groups, each carrying its own share of the answered probability.
func TestDistinctRRsetsSeparateGroups(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"),
	))
	// Both RRsets share their first sorted rdata so grouping must consider
	// the full content, not just the first record.
	m.on("1.1.1.1", "www.example.com", dns.TypeA, answerMsg(
		aRR("www.example.com", "1.0.0.1"),
		aRR("www.example.com", "9.9.9.9"),
	))
	m.on("2.2.2.2", "www.example.com", dns.TypeA, answerMsg(
		aRR("www.example.com", "1.0.0.1"),
		aRR("www.example.com", "8.8.8.8"),
	))

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

	assertSumsToOne(t, root)
	stats := root.SummaryStats()
	if stats == nil {
		t.Fatal("expected summary stats after calculation")
	}
	if prob := stats.ByStatus[StatusAnswered]; math.Abs(prob-1.0) > 1e-9 {
		t.Errorf("answered summary prob = %v, want 1.0", prob)
	}
	if len(stats.Answers) != 2 {
		t.Fatalf("expected 2 distinct answered RRsets, got %d: %v", len(stats.Answers), stats.Answers)
	}
	if stats.Answers[0].Key == stats.Answers[1].Key {
		t.Errorf("answer groups share key %q, want distinct keys", stats.Answers[0].Key)
	}
	for i, ans := range stats.Answers {
		if math.Abs(ans.Prob-0.5) > 1e-9 {
			t.Errorf("answer group %d (%q) prob = %v, want 0.5", i, ans.Key, ans.Prob)
		}
	}
	// Answers are sorted by key: "1.0.0.1@@@8.8.8.8" then "1.0.0.1@@@9.9.9.9".
	wantRdata := []string{"8.8.8.8", "9.9.9.9"}
	for i, ans := range stats.Answers {
		if !strings.Contains(ans.Key, "1.0.0.1@@@"+wantRdata[i]) {
			t.Errorf("answer group %d key = %q, want it to contain %q", i, ans.Key, "1.0.0.1@@@"+wantRdata[i])
		}
		if len(ans.RRs) != 2 {
			t.Errorf("answer group %d RRs = %v, want 2 records", i, ans.RRs)
		}
	}
}

func TestServfailErrorLeaf(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, answerMsg(aRR("www.example.com", "9.9.9.9")))
	m.on("2.2.2.2", "www.example.com", dns.TypeA, rcodeMsg(dns.RcodeServerFailure))

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

	assertSumsToOne(t, root)
	errs := leavesByStatus(root, StatusError)
	if len(errs) != 1 {
		t.Fatalf("expected 1 error leaf, got %v", root.StatsList())
	}
	if errs[0].Response.DQ.ErrorMessage != "Server failure (SERVFAIL)" {
		t.Errorf("error message = %q", errs[0].Response.DQ.ErrorMessage)
	}
	if math.Abs(errs[0].Prob-0.5) > 1e-9 {
		t.Errorf("error prob = %v, want 0.5", errs[0].Prob)
	}

	stats := root.SummaryStats()
	if math.Abs(stats.ByStatus[StatusError]-0.5) > 1e-9 ||
		math.Abs(stats.ByStatus[StatusAnswered]-0.5) > 1e-9 {
		t.Errorf("summary by status = %v", stats.ByStatus)
	}
	total := 0.0
	for _, prob := range stats.ByStatus {
		total += prob
	}
	if math.Abs(total-1.0) > 1e-9 {
		t.Errorf("summary probabilities sum to %v, want 1.0", total)
	}
}

// TestResolveSubtreeLeavesExcluded asserts that resolve-subtree leaves (the
// A <servername> lookups for glueless NS) never reach the main aggregation:
// they surface only through server weights.
func TestResolveSubtreeLeavesExcluded(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")))
	m.on("198.41.0.4", "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")))

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

	assertSumsToOne(t, root)
	for _, leaf := range root.StatsList() {
		if leaf.Response.Qname == "ns.other.net" {
			t.Errorf("resolve-subtree leaf leaked into main aggregation: %s", leaf.Key)
		}
	}
	if prob := root.SummaryStats().ByStatus[StatusAnswered]; math.Abs(prob-1.0) > 1e-9 {
		t.Errorf("answered summary prob = %v, want 1.0", prob)
	}
}

// TestResolveFailurePseudoIPCarriesMass asserts that a failed glue resolution
// keeps its probability: the failure becomes a "key:" pseudo-IP whose mass
// surfaces in the main aggregation as the failing (resolve) query.
func TestResolveFailurePseudoIPCarriesMass(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")))
	// The resolve of A ns.other.net fails at the root: SERVFAIL.
	m.on("198.41.0.4", "ns.other.net", dns.TypeA, rcodeMsg(dns.RcodeServerFailure))

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

	assertSumsToOne(t, root)
	errs := leavesByStatus(root, StatusError)
	if len(errs) != 1 {
		t.Fatalf("expected 1 error leaf from the failed resolve, got %v", root.StatsList())
	}
	leaf := errs[0]
	if math.Abs(leaf.Prob-0.5) > 1e-9 {
		t.Errorf("failed-resolve prob = %v, want 0.5", leaf.Prob)
	}
	if leaf.Response.Qname != "ns.other.net" {
		t.Errorf("failed-resolve leaf qname = %q, want the resolve target", leaf.Response.Qname)
	}
	if !strings.HasPrefix(leaf.Key, "key:error:") {
		t.Errorf("failed-resolve key = %q", leaf.Key)
	}
	// The leaf's referral is the resolve-subtree node; its parent is the
	// glueless referral, which carries the mass as a pseudo-IP server entry.
	glueless := leaf.Referral.Parent
	if glueless.Server != "ns.other.net" {
		t.Fatalf("glueless referral server = %q", glueless.Server)
	}
	hasPseudo := false
	for ip, weight := range glueless.ServerWeights {
		if strings.HasPrefix(ip, "key:") && math.Abs(weight-1.0) <= 1e-9 {
			hasPseudo = true
		}
	}
	if !hasPseudo {
		t.Errorf("expected a key: pseudo-IP with weight 1.0, got %v", glueless.ServerWeights)
	}
}

func TestSummaryStatsNilAndMemoised(t *testing.T) {
	var nilRef *Referral
	if nilRef.SummaryStats() != nil {
		t.Error("nil referral should produce nil summary")
	}
	uncalculated := newTestReferral("ns1.example.com", []string{"1.1.1.1"})
	if uncalculated.SummaryStats() != nil {
		t.Error("uncalculated referral should produce nil summary")
	}

	m := mockSimpleDelegation()
	_, root := runTraversal(t, testConfig(false), m, "www.example.com")
	first := root.SummaryStats()
	if first == nil {
		t.Fatal("expected summary stats")
	}
	if root.SummaryStats() != first {
		t.Error("summary stats should be memoised")
	}
}