package dns_test

import (
	"context"
	"net"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
	logtest "github.com/sirupsen/logrus/hooks/test"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	nbdns "github.com/netbirdio/netbird/client/internal/dns"
	"github.com/netbirdio/netbird/client/internal/dns/resutil"
	"github.com/netbirdio/netbird/client/internal/dns/test"
)

// TestHandlerChain_ServeDNS_Priorities tests that handlers are executed in priority order
func TestHandlerChain_ServeDNS_Priorities(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	// Create mock handlers for different priorities
	defaultHandler := &nbdns.MockHandler{}
	matchDomainHandler := &nbdns.MockHandler{}
	dnsRouteHandler := &nbdns.MockHandler{}

	// Setup handlers with different priorities
	chain.AddHandler("example.com.", defaultHandler, nbdns.PriorityDefault)
	chain.AddHandler("example.com.", matchDomainHandler, nbdns.PriorityUpstream)
	chain.AddHandler("example.com.", dnsRouteHandler, nbdns.PriorityDNSRoute)

	// Create test request
	r := new(dns.Msg)
	r.SetQuestion("example.com.", dns.TypeA)

	// Create test writer
	w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}

	// Setup expectations - only highest priority handler should be called
	dnsRouteHandler.On("ServeDNS", mock.Anything, r).Once()
	matchDomainHandler.On("ServeDNS", mock.Anything, r).Maybe()
	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe()

	// Execute
	chain.ServeDNS(w, r)

	// Verify all expectations were met
	dnsRouteHandler.AssertExpectations(t)
	matchDomainHandler.AssertExpectations(t)
	defaultHandler.AssertExpectations(t)
}

// TestHandlerChain_ServeDNS_DomainMatching tests various domain matching scenarios
func TestHandlerChain_ServeDNS_DomainMatching(t *testing.T) {
	tests := []struct {
		name            string
		handlerDomain   string
		queryDomain     string
		isWildcard      bool
		matchSubdomains bool
		shouldMatch     bool
	}{
		{
			name:            "exact match",
			handlerDomain:   "example.com.",
			queryDomain:     "example.com.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "subdomain with non-wildcard and MatchSubdomains true",
			handlerDomain:   "example.com.",
			queryDomain:     "sub.example.com.",
			isWildcard:      false,
			matchSubdomains: true,
			shouldMatch:     true,
		},
		{
			name:            "subdomain with non-wildcard and MatchSubdomains false",
			handlerDomain:   "example.com.",
			queryDomain:     "sub.example.com.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     false,
		},
		{
			name:            "wildcard match",
			handlerDomain:   "*.example.com.",
			queryDomain:     "sub.example.com.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "wildcard no match on apex",
			handlerDomain:   "*.example.com.",
			queryDomain:     "example.com.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     false,
		},
		{
			name:            "root zone match",
			handlerDomain:   ".",
			queryDomain:     "anything.com.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "no match different domain",
			handlerDomain:   "example.com.",
			queryDomain:     "example.org.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     false,
		},
		{
			name:            "single letter TLD exact match",
			handlerDomain:   "example.x.",
			queryDomain:     "example.x.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "single letter TLD subdomain match",
			handlerDomain:   "example.x.",
			queryDomain:     "sub.example.x.",
			isWildcard:      false,
			matchSubdomains: true,
			shouldMatch:     true,
		},
		{
			name:            "single letter TLD wildcard match",
			handlerDomain:   "*.example.x.",
			queryDomain:     "sub.example.x.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "two letter domain labels",
			handlerDomain:   "a.b.",
			queryDomain:     "a.b.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "single character domain",
			handlerDomain:   "x.",
			queryDomain:     "x.",
			isWildcard:      false,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "single character domain with subdomain match",
			handlerDomain:   "x.",
			queryDomain:     "sub.x.",
			isWildcard:      false,
			matchSubdomains: true,
			shouldMatch:     true,
		},
		{
			name:            "wildcard label-boundary mismatch (suffix overlap)",
			handlerDomain:   "*.b.test.",
			queryDomain:     "x.ab.test.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     false,
		},
		{
			name:            "wildcard label-boundary match",
			handlerDomain:   "*.b.test.",
			queryDomain:     "x.b.test.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "wildcard multi-label match",
			handlerDomain:   "*.b.test.",
			queryDomain:     "x.y.b.test.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     true,
		},
		{
			name:            "wildcard no match on multi-label apex",
			handlerDomain:   "*.b.test.",
			queryDomain:     "b.test.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     false,
		},
		{
			name:            "wildcard no match on unrelated suffix containment",
			handlerDomain:   "*.example.com.",
			queryDomain:     "notexample.com.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     false,
		},
		{
			name:            "wildcard accepts pattern registered without trailing dot",
			handlerDomain:   "*.b.test",
			queryDomain:     "x.b.test.",
			isWildcard:      true,
			matchSubdomains: false,
			shouldMatch:     true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chain := nbdns.NewHandlerChain()
			var handler dns.Handler

			if tt.matchSubdomains {
				mockSubHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
				handler = mockSubHandler
				if tt.shouldMatch {
					mockSubHandler.On("ServeDNS", mock.Anything, mock.Anything).Once()
				}
			} else {
				mockHandler := &nbdns.MockHandler{}
				handler = mockHandler
				if tt.shouldMatch {
					mockHandler.On("ServeDNS", mock.Anything, mock.Anything).Once()
				}
			}

			pattern := tt.handlerDomain
			if tt.isWildcard {
				pattern = "*." + tt.handlerDomain[2:]
			}

			chain.AddHandler(pattern, handler, nbdns.PriorityDefault)

			r := new(dns.Msg)
			r.SetQuestion(tt.queryDomain, dns.TypeA)
			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}

			chain.ServeDNS(w, r)

			if h, ok := handler.(*nbdns.MockHandler); ok {
				h.AssertExpectations(t)
			} else if h, ok := handler.(*nbdns.MockSubdomainHandler); ok {
				h.AssertExpectations(t)
			}
		})
	}
}

// TestHandlerChain_ServeDNS_OverlappingDomains tests behavior with overlapping domain patterns
func TestHandlerChain_ServeDNS_OverlappingDomains(t *testing.T) {
	tests := []struct {
		name     string
		handlers []struct {
			pattern  string
			priority int
		}
		queryDomain     string
		expectedCalls   int
		expectedHandler int // index of the handler that should be called
	}{
		{
			name: "wildcard and exact same priority - exact should win",
			handlers: []struct {
				pattern  string
				priority int
			}{
				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
				{pattern: "example.com.", priority: nbdns.PriorityDefault},
			},
			queryDomain:     "example.com.",
			expectedCalls:   1,
			expectedHandler: 1, // exact match handler should be called
		},
		{
			name: "higher priority wildcard over lower priority exact",
			handlers: []struct {
				pattern  string
				priority int
			}{
				{pattern: "example.com.", priority: nbdns.PriorityDefault},
				{pattern: "*.example.com.", priority: nbdns.PriorityDNSRoute},
			},
			queryDomain:     "test.example.com.",
			expectedCalls:   1,
			expectedHandler: 1, // higher priority wildcard handler should be called
		},
		{
			name: "multiple wildcards different priorities",
			handlers: []struct {
				pattern  string
				priority int
			}{
				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
				{pattern: "*.example.com.", priority: nbdns.PriorityUpstream},
				{pattern: "*.example.com.", priority: nbdns.PriorityDNSRoute},
			},
			queryDomain:     "test.example.com.",
			expectedCalls:   1,
			expectedHandler: 2, // highest priority handler should be called
		},
		{
			name: "subdomain with mix of patterns",
			handlers: []struct {
				pattern  string
				priority int
			}{
				{pattern: "*.example.com.", priority: nbdns.PriorityDefault},
				{pattern: "test.example.com.", priority: nbdns.PriorityUpstream},
				{pattern: "*.test.example.com.", priority: nbdns.PriorityDNSRoute},
			},
			queryDomain:     "sub.test.example.com.",
			expectedCalls:   1,
			expectedHandler: 2, // highest priority matching handler should be called
		},
		{
			name: "overlapping wildcard suffixes route to correct handler",
			handlers: []struct {
				pattern  string
				priority int
			}{
				{pattern: "*.b.test.", priority: nbdns.PriorityDNSRoute},
				{pattern: "*.ab.test.", priority: nbdns.PriorityDNSRoute},
			},
			queryDomain:     "app.ab.test.",
			expectedCalls:   1,
			expectedHandler: 1,
		},
		{
			name: "root zone with specific domain",
			handlers: []struct {
				pattern  string
				priority int
			}{
				{pattern: ".", priority: nbdns.PriorityDefault},
				{pattern: "example.com.", priority: nbdns.PriorityDNSRoute},
			},
			queryDomain:     "example.com.",
			expectedCalls:   1,
			expectedHandler: 1, // higher priority specific domain should win over root
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chain := nbdns.NewHandlerChain()
			var handlers []*nbdns.MockHandler

			// Setup handlers and expectations
			for i := range tt.handlers {
				handler := &nbdns.MockHandler{}
				handlers = append(handlers, handler)

				// Set expectation based on whether this handler should be called
				if i == tt.expectedHandler {
					handler.On("ServeDNS", mock.Anything, mock.Anything).Once()
				} else {
					handler.On("ServeDNS", mock.Anything, mock.Anything).Maybe()
				}

				chain.AddHandler(tt.handlers[i].pattern, handler, tt.handlers[i].priority)
			}

			// Create and execute request
			r := new(dns.Msg)
			r.SetQuestion(tt.queryDomain, dns.TypeA)
			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
			chain.ServeDNS(w, r)

			// Verify expectations
			for _, handler := range handlers {
				handler.AssertExpectations(t)
			}
		})
	}
}

// TestHandlerChain_ServeDNS_ChainContinuation tests the chain continuation functionality
func TestHandlerChain_ServeDNS_ChainContinuation(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	// Create handlers
	handler1 := &nbdns.MockHandler{}
	handler2 := &nbdns.MockHandler{}
	handler3 := &nbdns.MockHandler{}

	// Add handlers in priority order
	chain.AddHandler("example.com.", handler1, nbdns.PriorityDNSRoute)
	chain.AddHandler("example.com.", handler2, nbdns.PriorityUpstream)
	chain.AddHandler("example.com.", handler3, nbdns.PriorityDefault)

	// Create test request
	r := new(dns.Msg)
	r.SetQuestion("example.com.", dns.TypeA)

	// Setup mock responses to simulate chain continuation
	handler1.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
		// First handler signals continue
		w := args.Get(0).(*nbdns.ResponseWriterChain)
		resp := new(dns.Msg)
		resp.SetRcode(r, dns.RcodeNameError)
		resp.MsgHdr.Zero = true // Signal to continue
		assert.NoError(t, w.WriteMsg(resp))
	}).Once()

	handler2.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
		// Second handler signals continue
		w := args.Get(0).(*nbdns.ResponseWriterChain)
		resp := new(dns.Msg)
		resp.SetRcode(r, dns.RcodeNameError)
		resp.MsgHdr.Zero = true
		assert.NoError(t, w.WriteMsg(resp))
	}).Once()

	handler3.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
		// Last handler responds normally
		w := args.Get(0).(*nbdns.ResponseWriterChain)
		resp := new(dns.Msg)
		resp.SetRcode(r, dns.RcodeSuccess)
		assert.NoError(t, w.WriteMsg(resp))
	}).Once()

	// Execute
	w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
	chain.ServeDNS(w, r)

	// Verify all handlers were called in order
	handler1.AssertExpectations(t)
	handler2.AssertExpectations(t)
	handler3.AssertExpectations(t)
}

func TestHandlerChain_PriorityDeregistration(t *testing.T) {
	tests := []struct {
		name string
		ops  []struct {
			action   string // "add" or "remove"
			pattern  string
			priority int
		}
		query         string
		expectedCalls map[int]bool // map[priority]shouldBeCalled
	}{
		{
			name: "remove high priority keeps lower priority handler",
			ops: []struct {
				action   string
				pattern  string
				priority int
			}{
				{"add", "example.com.", nbdns.PriorityDNSRoute},
				{"add", "example.com.", nbdns.PriorityUpstream},
				{"remove", "example.com.", nbdns.PriorityDNSRoute},
			},
			query: "example.com.",
			expectedCalls: map[int]bool{
				nbdns.PriorityDNSRoute: false,
				nbdns.PriorityUpstream: true,
			},
		},
		{
			name: "remove lower priority keeps high priority handler",
			ops: []struct {
				action   string
				pattern  string
				priority int
			}{
				{"add", "example.com.", nbdns.PriorityDNSRoute},
				{"add", "example.com.", nbdns.PriorityUpstream},
				{"remove", "example.com.", nbdns.PriorityUpstream},
			},
			query: "example.com.",
			expectedCalls: map[int]bool{
				nbdns.PriorityDNSRoute: true,
				nbdns.PriorityUpstream: false,
			},
		},
		{
			name: "remove all handlers in order",
			ops: []struct {
				action   string
				pattern  string
				priority int
			}{
				{"add", "example.com.", nbdns.PriorityDNSRoute},
				{"add", "example.com.", nbdns.PriorityUpstream},
				{"add", "example.com.", nbdns.PriorityDefault},
				{"remove", "example.com.", nbdns.PriorityDNSRoute},
				{"remove", "example.com.", nbdns.PriorityUpstream},
			},
			query: "example.com.",
			expectedCalls: map[int]bool{
				nbdns.PriorityDNSRoute: false,
				nbdns.PriorityUpstream: false,
				nbdns.PriorityDefault:  true,
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chain := nbdns.NewHandlerChain()
			handlers := make(map[int]*nbdns.MockHandler)

			// Execute operations
			for _, op := range tt.ops {
				if op.action == "add" {
					handler := &nbdns.MockHandler{}
					handlers[op.priority] = handler
					chain.AddHandler(op.pattern, handler, op.priority)
				} else {
					chain.RemoveHandler(op.pattern, op.priority)
				}
			}

			// Create test request
			r := new(dns.Msg)
			r.SetQuestion(tt.query, dns.TypeA)
			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}

			// Setup expectations
			for priority, handler := range handlers {
				if shouldCall, exists := tt.expectedCalls[priority]; exists && shouldCall {
					handler.On("ServeDNS", mock.Anything, r).Once()
				} else {
					handler.On("ServeDNS", mock.Anything, r).Maybe()
				}
			}

			// Execute request
			chain.ServeDNS(w, r)

			// Verify expectations
			for _, handler := range handlers {
				handler.AssertExpectations(t)
			}
		})
	}
}

func TestHandlerChain_MultiPriorityHandling(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	testDomain := "example.com."
	testQuery := "test.example.com."

	// Create handlers with MatchSubdomains enabled
	routeHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
	matchHandler := &nbdns.MockSubdomainHandler{Subdomains: true}
	defaultHandler := &nbdns.MockSubdomainHandler{Subdomains: true}

	// Create test request that will be reused
	r := new(dns.Msg)
	r.SetQuestion(testQuery, dns.TypeA)

	// Keep track of mocks for the final assertion in Step 4
	mocks := []*nbdns.MockSubdomainHandler{routeHandler, matchHandler, defaultHandler}

	// Add handlers in mixed order
	chain.AddHandler(testDomain, defaultHandler, nbdns.PriorityDefault)
	chain.AddHandler(testDomain, routeHandler, nbdns.PriorityDNSRoute)
	chain.AddHandler(testDomain, matchHandler, nbdns.PriorityUpstream)

	// Test 1: Initial state
	w1 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
	// Highest priority handler (routeHandler) should be called
	routeHandler.On("ServeDNS", mock.Anything, r).Return().Once()
	matchHandler.On("ServeDNS", mock.Anything, r).Maybe()   // Ensure others are not expected yet
	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe() // Ensure others are not expected yet

	chain.ServeDNS(w1, r)
	routeHandler.AssertExpectations(t)

	routeHandler.ExpectedCalls = nil
	routeHandler.Calls = nil
	matchHandler.ExpectedCalls = nil
	matchHandler.Calls = nil
	defaultHandler.ExpectedCalls = nil
	defaultHandler.Calls = nil

	// Test 2: Remove highest priority handler
	chain.RemoveHandler(testDomain, nbdns.PriorityDNSRoute)

	w2 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
	// Now middle priority handler (matchHandler) should be called
	matchHandler.On("ServeDNS", mock.Anything, r).Return().Once()
	defaultHandler.On("ServeDNS", mock.Anything, r).Maybe() // Ensure default is not expected yet

	chain.ServeDNS(w2, r)
	matchHandler.AssertExpectations(t)

	matchHandler.ExpectedCalls = nil
	matchHandler.Calls = nil
	defaultHandler.ExpectedCalls = nil
	defaultHandler.Calls = nil

	// Test 3: Remove middle priority handler
	chain.RemoveHandler(testDomain, nbdns.PriorityUpstream)

	w3 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
	// Now lowest priority handler (defaultHandler) should be called
	defaultHandler.On("ServeDNS", mock.Anything, r).Return().Once()

	chain.ServeDNS(w3, r)
	defaultHandler.AssertExpectations(t)

	defaultHandler.ExpectedCalls = nil
	defaultHandler.Calls = nil

	// Test 4: Remove last handler
	chain.RemoveHandler(testDomain, nbdns.PriorityDefault)

	w4 := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}
	chain.ServeDNS(w4, r) // Call ServeDNS on the now empty chain for this domain

	for _, m := range mocks {
		m.AssertNumberOfCalls(t, "ServeDNS", 0)
	}
}

func TestHandlerChain_CaseSensitivity(t *testing.T) {
	tests := []struct {
		name        string
		scenario    string
		addHandlers []struct {
			pattern     string
			priority    int
			subdomains  bool
			shouldMatch bool
		}
		query         string
		expectedCalls int
	}{
		{
			name:     "case insensitive exact match",
			scenario: "handler registered lowercase, query uppercase",
			addHandlers: []struct {
				pattern     string
				priority    int
				subdomains  bool
				shouldMatch bool
			}{
				{"example.com.", nbdns.PriorityDefault, false, true},
			},
			query:         "EXAMPLE.COM.",
			expectedCalls: 1,
		},
		{
			name:     "case insensitive wildcard match",
			scenario: "handler registered mixed case wildcard, query different case",
			addHandlers: []struct {
				pattern     string
				priority    int
				subdomains  bool
				shouldMatch bool
			}{
				{"*.Example.Com.", nbdns.PriorityDefault, false, true},
			},
			query:         "sub.EXAMPLE.COM.",
			expectedCalls: 1,
		},
		{
			name:     "multiple handlers different case same domain",
			scenario: "second handler should replace first despite case difference",
			addHandlers: []struct {
				pattern     string
				priority    int
				subdomains  bool
				shouldMatch bool
			}{
				{"EXAMPLE.COM.", nbdns.PriorityDefault, false, false},
				{"example.com.", nbdns.PriorityDefault, false, true},
			},
			query:         "ExAmPlE.cOm.",
			expectedCalls: 1,
		},
		{
			name:     "subdomain matching case insensitive",
			scenario: "handler with MatchSubdomains true should match regardless of case",
			addHandlers: []struct {
				pattern     string
				priority    int
				subdomains  bool
				shouldMatch bool
			}{
				{"example.com.", nbdns.PriorityDefault, true, true},
			},
			query:         "SUB.EXAMPLE.COM.",
			expectedCalls: 1,
		},
		{
			name:     "root zone case insensitive",
			scenario: "root zone handler should match regardless of case",
			addHandlers: []struct {
				pattern     string
				priority    int
				subdomains  bool
				shouldMatch bool
			}{
				{".", nbdns.PriorityDefault, false, true},
			},
			query:         "EXAMPLE.COM.",
			expectedCalls: 1,
		},
		{
			name:     "multiple handlers different priority",
			scenario: "should call higher priority handler despite case differences",
			addHandlers: []struct {
				pattern     string
				priority    int
				subdomains  bool
				shouldMatch bool
			}{
				{"EXAMPLE.COM.", nbdns.PriorityDefault, false, false},
				{"example.com.", nbdns.PriorityUpstream, false, false},
				{"Example.Com.", nbdns.PriorityDNSRoute, false, true},
			},
			query:         "example.com.",
			expectedCalls: 1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chain := nbdns.NewHandlerChain()
			handlerCalls := make(map[string]bool) // track which patterns were called

			// Add handlers according to test case
			for _, h := range tt.addHandlers {
				var handler dns.Handler
				pattern := h.pattern // capture pattern for closure

				if h.subdomains {
					subHandler := &nbdns.MockSubdomainHandler{
						Subdomains: true,
					}
					if h.shouldMatch {
						subHandler.On("ServeDNS", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
							handlerCalls[pattern] = true
							w := args.Get(0).(dns.ResponseWriter)
							r := args.Get(1).(*dns.Msg)
							resp := new(dns.Msg)
							resp.SetRcode(r, dns.RcodeSuccess)
							assert.NoError(t, w.WriteMsg(resp))
						}).Once()
					}
					handler = subHandler
				} else {
					mockHandler := &nbdns.MockHandler{}
					if h.shouldMatch {
						mockHandler.On("ServeDNS", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
							handlerCalls[pattern] = true
							w := args.Get(0).(dns.ResponseWriter)
							r := args.Get(1).(*dns.Msg)
							resp := new(dns.Msg)
							resp.SetRcode(r, dns.RcodeSuccess)
							assert.NoError(t, w.WriteMsg(resp))
						}).Once()
					}
					handler = mockHandler
				}

				chain.AddHandler(pattern, handler, h.priority)
			}

			// Execute request
			r := new(dns.Msg)
			r.SetQuestion(tt.query, dns.TypeA)
			chain.ServeDNS(&test.MockResponseWriter{}, r)

			// Verify each handler was called exactly as expected
			for _, h := range tt.addHandlers {
				wasCalled := handlerCalls[h.pattern]
				assert.Equal(t, h.shouldMatch, wasCalled,
					"Handler for pattern %q was %s when it should%s have been",
					h.pattern,
					map[bool]string{true: "called", false: "not called"}[wasCalled],
					map[bool]string{true: "", false: " not"}[wasCalled == h.shouldMatch])
			}

			// Verify total number of calls
			assert.Equal(t, tt.expectedCalls, len(handlerCalls),
				"Wrong number of total handler calls")
		})
	}
}

func TestHandlerChain_DomainSpecificityOrdering(t *testing.T) {
	tests := []struct {
		name     string
		scenario string
		ops      []struct {
			action    string
			pattern   string
			priority  int
			subdomain bool
		}
		query         string
		expectedMatch string
	}{
		{
			name:     "more specific domain matches first",
			scenario: "sub.example.com should match before example.com",
			ops: []struct {
				action    string
				pattern   string
				priority  int
				subdomain bool
			}{
				{"add", "example.com.", nbdns.PriorityUpstream, true},
				{"add", "sub.example.com.", nbdns.PriorityUpstream, false},
			},
			query:         "sub.example.com.",
			expectedMatch: "sub.example.com.",
		},
		{
			name:     "more specific domain matches first, both match subdomains",
			scenario: "sub.example.com should match before example.com",
			ops: []struct {
				action    string
				pattern   string
				priority  int
				subdomain bool
			}{
				{"add", "example.com.", nbdns.PriorityUpstream, true},
				{"add", "sub.example.com.", nbdns.PriorityUpstream, true},
			},
			query:         "sub.example.com.",
			expectedMatch: "sub.example.com.",
		},
		{
			name:     "maintain specificity order after removal",
			scenario: "after removing most specific, should fall back to less specific",
			ops: []struct {
				action    string
				pattern   string
				priority  int
				subdomain bool
			}{
				{"add", "example.com.", nbdns.PriorityUpstream, true},
				{"add", "sub.example.com.", nbdns.PriorityUpstream, true},
				{"add", "test.sub.example.com.", nbdns.PriorityUpstream, false},
				{"remove", "test.sub.example.com.", nbdns.PriorityUpstream, false},
			},
			query:         "test.sub.example.com.",
			expectedMatch: "sub.example.com.",
		},
		{
			name:     "priority overrides specificity",
			scenario: "less specific domain with higher priority should match first",
			ops: []struct {
				action    string
				pattern   string
				priority  int
				subdomain bool
			}{
				{"add", "sub.example.com.", nbdns.PriorityUpstream, false},
				{"add", "example.com.", nbdns.PriorityDNSRoute, true},
			},
			query:         "sub.example.com.",
			expectedMatch: "example.com.",
		},
		{
			name:     "equal priority respects specificity",
			scenario: "with equal priority, more specific domain should match",
			ops: []struct {
				action    string
				pattern   string
				priority  int
				subdomain bool
			}{
				{"add", "example.com.", nbdns.PriorityUpstream, true},
				{"add", "other.example.com.", nbdns.PriorityUpstream, true},
				{"add", "sub.example.com.", nbdns.PriorityUpstream, false},
			},
			query:         "sub.example.com.",
			expectedMatch: "sub.example.com.",
		},
		{
			name:     "specific matches before wildcard",
			scenario: "specific domain should match before wildcard at same priority",
			ops: []struct {
				action    string
				pattern   string
				priority  int
				subdomain bool
			}{
				{"add", "*.example.com.", nbdns.PriorityDNSRoute, false},
				{"add", "sub.example.com.", nbdns.PriorityDNSRoute, false},
			},
			query:         "sub.example.com.",
			expectedMatch: "sub.example.com.",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chain := nbdns.NewHandlerChain()
			handlers := make(map[string]*nbdns.MockSubdomainHandler)

			for _, op := range tt.ops {
				if op.action == "add" {
					handler := &nbdns.MockSubdomainHandler{Subdomains: op.subdomain}
					handlers[op.pattern] = handler
					chain.AddHandler(op.pattern, handler, op.priority)
				} else {
					chain.RemoveHandler(op.pattern, op.priority)
				}
			}

			r := new(dns.Msg)
			r.SetQuestion(tt.query, dns.TypeA)
			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}

			// Setup handler expectations
			for pattern, handler := range handlers {
				if pattern == tt.expectedMatch {
					handler.On("ServeDNS", mock.Anything, r).Run(func(args mock.Arguments) {
						w := args.Get(0).(dns.ResponseWriter)
						r := args.Get(1).(*dns.Msg)
						resp := new(dns.Msg)
						resp.SetReply(r)
						assert.NoError(t, w.WriteMsg(resp))
					}).Once()
				}
			}

			chain.ServeDNS(w, r)

			for pattern, handler := range handlers {
				if pattern == tt.expectedMatch {
					handler.AssertNumberOfCalls(t, "ServeDNS", 1)
				} else {
					handler.AssertNumberOfCalls(t, "ServeDNS", 0)
				}
			}
		})
	}
}

func TestHandlerChain_AddRemoveRoundtrip(t *testing.T) {
	tests := []struct {
		name            string
		addPattern      string
		removePattern   string
		queryPattern    string
		shouldBeRemoved bool
		description     string
	}{
		{
			name:            "exact same pattern",
			addPattern:      "example.com.",
			removePattern:   "example.com.",
			queryPattern:    "example.com.",
			shouldBeRemoved: true,
			description:     "Adding and removing with identical patterns",
		},
		{
			name:            "case difference",
			addPattern:      "Example.Com.",
			removePattern:   "EXAMPLE.COM.",
			queryPattern:    "example.com.",
			shouldBeRemoved: true,
			description:     "Adding with mixed case, removing with uppercase",
		},
		{
			name:            "reversed case difference",
			addPattern:      "EXAMPLE.ORG.",
			removePattern:   "example.org.",
			queryPattern:    "example.org.",
			shouldBeRemoved: true,
			description:     "Adding with uppercase, removing with lowercase",
		},
		{
			name:            "add wildcard, remove wildcard",
			addPattern:      "*.example.com.",
			removePattern:   "*.example.com.",
			queryPattern:    "sub.example.com.",
			shouldBeRemoved: true,
			description:     "Adding and removing with identical wildcard patterns",
		},
		{
			name:            "add wildcard, remove transformed pattern",
			addPattern:      "*.example.net.",
			removePattern:   "example.net.",
			queryPattern:    "sub.example.net.",
			shouldBeRemoved: false,
			description:     "Adding with wildcard, removing with non-wildcard pattern",
		},
		{
			name:            "add transformed pattern, remove wildcard",
			addPattern:      "example.io.",
			removePattern:   "*.example.io.",
			queryPattern:    "example.io.",
			shouldBeRemoved: false,
			description:     "Adding with non-wildcard pattern, removing with wildcard pattern",
		},
		{
			name:            "trailing dot difference",
			addPattern:      "example.dev",
			removePattern:   "example.dev.",
			queryPattern:    "example.dev.",
			shouldBeRemoved: true,
			description:     "Adding without trailing dot, removing with trailing dot",
		},
		{
			name:            "reversed trailing dot difference",
			addPattern:      "example.app.",
			removePattern:   "example.app",
			queryPattern:    "example.app.",
			shouldBeRemoved: true,
			description:     "Adding with trailing dot, removing without trailing dot",
		},
		{
			name:            "mixed case and wildcard",
			addPattern:      "*.Example.Site.",
			removePattern:   "*.EXAMPLE.SITE.",
			queryPattern:    "sub.example.site.",
			shouldBeRemoved: true,
			description:     "Adding mixed case wildcard, removing uppercase wildcard",
		},
		{
			name:            "root zone",
			addPattern:      ".",
			removePattern:   ".",
			queryPattern:    "random.domain.",
			shouldBeRemoved: true,
			description:     "Adding and removing root zone",
		},
		{
			name:            "wrong domain",
			addPattern:      "example.com.",
			removePattern:   "different.com.",
			queryPattern:    "example.com.",
			shouldBeRemoved: false,
			description:     "Adding one domain, trying to remove a different domain",
		},
		{
			name:            "subdomain mismatch",
			addPattern:      "sub.example.com.",
			removePattern:   "example.com.",
			queryPattern:    "sub.example.com.",
			shouldBeRemoved: false,
			description:     "Adding subdomain, trying to remove parent domain",
		},
		{
			name:            "parent domain mismatch",
			addPattern:      "example.com.",
			removePattern:   "sub.example.com.",
			queryPattern:    "example.com.",
			shouldBeRemoved: false,
			description:     "Adding parent domain, trying to remove subdomain",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chain := nbdns.NewHandlerChain()

			handler := &nbdns.MockHandler{}
			r := new(dns.Msg)
			r.SetQuestion(tt.queryPattern, dns.TypeA)
			w := &nbdns.ResponseWriterChain{ResponseWriter: &test.MockResponseWriter{}}

			// First verify no handler is called before adding any
			chain.ServeDNS(w, r)
			handler.AssertNotCalled(t, "ServeDNS")

			// Add handler
			chain.AddHandler(tt.addPattern, handler, nbdns.PriorityDefault)

			// Verify handler is called after adding
			handler.On("ServeDNS", mock.Anything, r).Once()
			chain.ServeDNS(w, r)
			handler.AssertExpectations(t)

			// Reset mock for the next test
			handler.ExpectedCalls = nil

			// Remove handler
			chain.RemoveHandler(tt.removePattern, nbdns.PriorityDefault)

			// Set up expectations based on whether removal should succeed
			if !tt.shouldBeRemoved {
				handler.On("ServeDNS", mock.Anything, r).Once()
			}

			// Test if handler is still called after removal attempt
			chain.ServeDNS(w, r)

			if tt.shouldBeRemoved {
				handler.AssertNotCalled(t, "ServeDNS",
					"Handler should not be called after successful removal with pattern %q",
					tt.removePattern)
			} else {
				handler.AssertExpectations(t)
				handler.ExpectedCalls = nil
			}
		})
	}
}

// answeringHandler writes a fixed A record to ack the query. Used to verify
// which handler ResolveInternal dispatches to.
type answeringHandler struct {
	name string
	ip   string
}

func (h *answeringHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	resp := &dns.Msg{}
	resp.SetReply(r)
	resp.Answer = []dns.RR{&dns.A{
		Hdr: dns.RR_Header{Name: r.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
		A:   net.ParseIP(h.ip).To4(),
	}}
	_ = w.WriteMsg(resp)
}

func (h *answeringHandler) String() string { return h.name }

func TestHandlerChain_ResolveInternal_SkipsAboveMaxPriority(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	high := &answeringHandler{name: "high", ip: "10.0.0.1"}
	low := &answeringHandler{name: "low", ip: "10.0.0.2"}

	chain.AddHandler("example.com.", high, nbdns.PriorityMgmtCache)
	chain.AddHandler("example.com.", low, nbdns.PriorityUpstream)

	r := new(dns.Msg)
	r.SetQuestion("example.com.", dns.TypeA)

	resp, err := chain.ResolveInternal(context.Background(), r, nbdns.PriorityUpstream)
	assert.NoError(t, err)
	assert.NotNil(t, resp)
	assert.Equal(t, 1, len(resp.Answer))
	a, ok := resp.Answer[0].(*dns.A)
	assert.True(t, ok)
	assert.Equal(t, "10.0.0.2", a.A.String(), "should skip mgmtCache handler and resolve via upstream")
}

func TestHandlerChain_ResolveInternal_ErrorWhenNoMatch(t *testing.T) {
	chain := nbdns.NewHandlerChain()
	high := &answeringHandler{name: "high", ip: "10.0.0.1"}
	chain.AddHandler("example.com.", high, nbdns.PriorityMgmtCache)

	r := new(dns.Msg)
	r.SetQuestion("example.com.", dns.TypeA)

	_, err := chain.ResolveInternal(context.Background(), r, nbdns.PriorityUpstream)
	assert.Error(t, err, "no handler at or below maxPriority should error")
}

// rawWriteHandler packs a response and calls ResponseWriter.Write directly
// (instead of WriteMsg), exercising the internalResponseWriter.Write path.
type rawWriteHandler struct {
	ip string
}

func (h *rawWriteHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	resp := &dns.Msg{}
	resp.SetReply(r)
	resp.Answer = []dns.RR{&dns.A{
		Hdr: dns.RR_Header{Name: r.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
		A:   net.ParseIP(h.ip).To4(),
	}}
	packed, err := resp.Pack()
	if err != nil {
		return
	}
	_, _ = w.Write(packed)
}

func TestHandlerChain_ResolveInternal_CapturesRawWrite(t *testing.T) {
	chain := nbdns.NewHandlerChain()
	chain.AddHandler("example.com.", &rawWriteHandler{ip: "10.0.0.3"}, nbdns.PriorityUpstream)

	r := new(dns.Msg)
	r.SetQuestion("example.com.", dns.TypeA)

	resp, err := chain.ResolveInternal(context.Background(), r, nbdns.PriorityUpstream)
	assert.NoError(t, err)
	require.NotNil(t, resp)
	require.Len(t, resp.Answer, 1)
	a, ok := resp.Answer[0].(*dns.A)
	require.True(t, ok)
	assert.Equal(t, "10.0.0.3", a.A.String(), "handlers calling Write(packed) must still surface their answer")
}

func TestHandlerChain_ResolveInternal_EmptyQuestion(t *testing.T) {
	chain := nbdns.NewHandlerChain()
	_, err := chain.ResolveInternal(context.Background(), new(dns.Msg), nbdns.PriorityUpstream)
	assert.Error(t, err)
}

// hangingHandler blocks indefinitely until closed, simulating a wedged upstream.
type hangingHandler struct {
	block chan struct{}
}

func (h *hangingHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	<-h.block
	resp := &dns.Msg{}
	resp.SetReply(r)
	_ = w.WriteMsg(resp)
}

func (h *hangingHandler) String() string { return "hangingHandler" }

func TestHandlerChain_ResolveInternal_HonorsContextTimeout(t *testing.T) {
	chain := nbdns.NewHandlerChain()
	h := &hangingHandler{block: make(chan struct{})}
	defer close(h.block)

	chain.AddHandler("example.com.", h, nbdns.PriorityUpstream)

	r := new(dns.Msg)
	r.SetQuestion("example.com.", dns.TypeA)

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	start := time.Now()
	_, err := chain.ResolveInternal(ctx, r, nbdns.PriorityUpstream)
	elapsed := time.Since(start)

	assert.Error(t, err)
	assert.ErrorIs(t, err, context.DeadlineExceeded)
	assert.Less(t, elapsed, 500*time.Millisecond, "ResolveInternal must return shortly after ctx deadline")
}

// requestSoftNegative asks the chain to soften a negative verdict produced by
// the handlers that run after the caller defers, reporting whether the chain
// supports the signal. Written as a type assertion so these tests compile
// against a chain that does not support it yet.
func requestSoftNegative(w dns.ResponseWriter) bool {
	sn, ok := w.(interface{ RequestSoftNegative() })
	if ok {
		sn.RequestSoftNegative()
	}
	return ok
}

// deferringHandler defers to the next handler in the chain, optionally asking
// for the negative verdict of whatever answers instead to be softened. This is
// what a DNS route handler does for a record type its routing peer cannot
// resolve.
type deferringHandler struct {
	softNegative bool
	called       bool
	supported    bool
}

func (h *deferringHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	h.called = true
	if h.softNegative {
		h.supported = requestSoftNegative(w)
		// The real handler records why it stepped aside, for the response log
		// line of whichever handler answers instead.
		resutil.SetMeta(w, resutil.MetaKeyDeferredBy, "test handler")
	}
	resp := new(dns.Msg)
	resp.SetRcode(r, dns.RcodeNameError)
	resp.MsgHdr.Zero = true
	_ = w.WriteMsg(resp)
}

// nxdomainHandler answers an authoritative NXDOMAIN with an SOA in the
// authority section, the way a public resolver answers for a name that only
// exists inside the routed network.
type nxdomainHandler struct{}

func (h *nxdomainHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	resp := new(dns.Msg)
	resp.SetRcode(r, dns.RcodeNameError)
	resp.Ns = []dns.RR{&dns.SOA{
		Hdr:     dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: 3600},
		Ns:      "ns1.example.com.",
		Mbox:    "hostmaster.example.com.",
		Minttl:  3600,
		Expire:  604800,
		Refresh: 7200,
		Retry:   3600,
	}}
	_ = w.WriteMsg(resp)
}

// TestHandlerChain_SoftNegative_DowngradesDownstreamNXDOMAIN is the whole point
// of the soft-negative signal: a route handler may only defer a query to the
// public chain if the answer cannot poison the routed name. NXDOMAIN is cached
// for the name and every type under it (RFC 2308, RFC 8020), so it has to be
// rewritten to NODATA, which is cached per name and type only.
func TestHandlerChain_SoftNegative_DowngradesDownstreamNXDOMAIN(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	route := &deferringHandler{softNegative: true}
	chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
	chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)

	r := new(dns.Msg)
	r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)

	mw := &test.MockResponseWriter{}
	chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

	resp := mw.GetLastResponse()
	require.NotNil(t, resp, "a response must reach the client")
	assert.True(t, route.called, "the route handler must run first")
	require.True(t, route.supported, "the chain writer must accept a soft-negative request")
	assert.Equal(t, dns.RcodeSuccess, resp.Rcode, "NXDOMAIN must be softened to NODATA")
	assert.Empty(t, resp.Answer, "a softened negative carries no answer")
	assert.Empty(t, resp.Ns, "the downstream zone's SOA must not set the negative TTL for a name we overrode")
}

// responseLineFor returns the chain's response log line for one query. Raising
// the level to trace also unmutes whatever else is logging in this package, so
// the line has to be picked by the name it was asked about rather than by being
// the last one seen.
func responseLineFor(hook *logtest.Hook, qname string) string {
	for _, e := range hook.AllEntries() {
		if strings.HasPrefix(e.Message, "response:") && strings.Contains(e.Message, qname) {
			return e.Message
		}
	}
	return ""
}

// TestHandlerChain_SoftNegative_IsVisibleToTheClient covers observability of the
// rewrite. The reply we hand the application travels over loopback, which the
// bundle capture does not see, so a softened verdict has to say so on the wire:
// without it an empty answer is indistinguishable from a real "no such record"
// in a dig output or a capture taken next to the application.
func TestHandlerChain_SoftNegative_IsVisibleToTheClient(t *testing.T) {
	// One hook for the whole test: logtest installs it on the standard logger
	// and logrus has no way to take it off again, so a hook per subtest would
	// leave several behind buffering every later line in the package.
	hook := logtest.NewGlobal()
	t.Cleanup(hook.Reset)

	newChain := func() *nbdns.HandlerChain {
		chain := nbdns.NewHandlerChain()
		chain.AddHandler("*.example.com.", &deferringHandler{softNegative: true}, nbdns.PriorityDNSRoute)
		chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)
		return chain
	}

	t.Run("EDNS0 client gets an extended error", func(t *testing.T) {
		r := new(dns.Msg)
		r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)
		r.SetEdns0(dns.DefaultMsgSize, false)

		mw := &test.MockResponseWriter{}
		newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

		resp := mw.GetLastResponse()
		require.NotNil(t, resp)
		require.Equal(t, dns.RcodeSuccess, resp.Rcode)

		ede, ok := resutil.ExtractEDE(resp)
		require.True(t, ok, "a softened verdict must carry an extended DNS error")
		assert.Equal(t, resutil.EDENetbirdSoftenedNegative, ede.InfoCode)
		assert.Contains(t, ede.ExtraText, "netbird", "the text must name us as the source of the rewrite")
	})

	t.Run("plain client gets no OPT", func(t *testing.T) {
		r := new(dns.Msg)
		r.SetQuestion("_mongodb._tcp.db.example.com.", dns.TypeSRV)

		mw := &test.MockResponseWriter{}
		newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

		resp := mw.GetLastResponse()
		require.NotNil(t, resp)
		assert.Nil(t, resp.IsEdns0(), "RFC 6891 forbids an OPT toward a client that did not advertise EDNS0")
	})

	// The response line is what support reads out of a debug bundle, and it now
	// carries several annotations at once. Built from a map, their order would
	// differ on every query, so the same event never looks the same twice.
	t.Run("log fields keep a stable order", func(t *testing.T) {
		const qname = "_mongodb._tcp.stable.example.com."

		prev := log.GetLevel()
		log.SetLevel(log.TraceLevel)
		t.Cleanup(func() { log.SetLevel(prev) })

		lineFor := func() string {
			hook.Reset()

			r := new(dns.Msg)
			r.SetQuestion(qname, dns.TypeSRV)

			mw := &test.MockResponseWriter{}
			newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

			line := responseLineFor(hook, qname)
			// The duration differs per query and is not what we compare.
			line, _, _ = strings.Cut(line, " took=")
			return line
		}

		first := lineFor()
		require.NotEmpty(t, first, "the chain must log the response it wrote")
		require.Contains(t, first, "deferred_by=", "the line must carry more than one annotation to be worth ordering")
		require.Contains(t, first, "softened=")

		for range 20 {
			assert.Equal(t, first, lineFor(), "the same event must produce the same line")
		}
	})

	t.Run("logged with the reason it was deferred", func(t *testing.T) {
		const qname = "_mongodb._tcp.reason.example.com."

		hook.Reset()

		prev := log.GetLevel()
		log.SetLevel(log.TraceLevel)
		t.Cleanup(func() { log.SetLevel(prev) })

		r := new(dns.Msg)
		r.SetQuestion(qname, dns.TypeSRV)

		mw := &test.MockResponseWriter{}
		newChain().ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

		response := responseLineFor(hook, qname)
		require.NotEmpty(t, response, "the chain must log the response it wrote")
		assert.Contains(t, response, "softened=", "the log must show the verdict was rewritten")
		assert.Contains(t, response, "deferred_by=", "the log must name the handler that deferred")
	})
}

// TestHandlerChain_SoftNegative_NoHandlerBelow covers a client with no primary
// nameserver group: the deferred query reaches the end of the chain unanswered.
// REFUSED would say the name is not served here while the route serves its
// addresses, and a stub that acts on that by asking elsewhere can bring back an
// NXDOMAIN for the whole name. The answer must be an empty, uncacheable NODATA.
func TestHandlerChain_SoftNegative_NoHandlerBelow(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	route := &deferringHandler{softNegative: true}
	chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)

	r := new(dns.Msg)
	r.SetQuestion("db.example.com.", dns.TypeHTTPS)

	mw := &test.MockResponseWriter{}
	chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

	resp := mw.GetLastResponse()
	require.NotNil(t, resp, "a response must reach the client")
	assert.Equal(t, dns.RcodeSuccess, resp.Rcode, "an unanswered soft-negative query must be NODATA, not REFUSED")
	assert.Empty(t, resp.Answer)
	assert.Empty(t, resp.Ns,
		"no SOA, so RFC 2308 keeps the empty answer out of negative caches and it cannot outlive the route")
}

// TestHandlerChain_SoftNegative_KeepsRealAnswers guards the other direction:
// softening applies to negative verdicts only. A real answer from a downstream
// handler must reach the client untouched.
func TestHandlerChain_SoftNegative_KeepsRealAnswers(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	route := &deferringHandler{softNegative: true}
	chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
	chain.AddHandler(".", &answeringHandler{name: "public", ip: "203.0.113.10"}, nbdns.PriorityDefault)

	r := new(dns.Msg)
	r.SetQuestion("db.example.com.", dns.TypeA)

	mw := &test.MockResponseWriter{}
	chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

	resp := mw.GetLastResponse()
	require.NotNil(t, resp, "a response must reach the client")
	assert.Equal(t, dns.RcodeSuccess, resp.Rcode)
	require.Len(t, resp.Answer, 1, "the downstream answer must pass through")
}

// TestHandlerChain_NXDOMAINPreservedWithoutSoftNegative makes sure the
// softening is opt-in: an ordinary chain continuation still yields NXDOMAIN, so
// genuine non-existence keeps being reported.
func TestHandlerChain_NXDOMAINPreservedWithoutSoftNegative(t *testing.T) {
	chain := nbdns.NewHandlerChain()

	route := &deferringHandler{}
	chain.AddHandler("*.example.com.", route, nbdns.PriorityDNSRoute)
	chain.AddHandler(".", &nxdomainHandler{}, nbdns.PriorityDefault)

	r := new(dns.Msg)
	r.SetQuestion("nope.example.com.", dns.TypeA)

	mw := &test.MockResponseWriter{}
	chain.ServeDNS(&nbdns.ResponseWriterChain{ResponseWriter: mw}, r)

	resp := mw.GetLastResponse()
	require.NotNil(t, resp, "a response must reach the client")
	assert.Equal(t, dns.RcodeNameError, resp.Rcode, "without the signal a real NXDOMAIN must survive")
}

func TestHandlerChain_HasRootHandlerAtOrBelow(t *testing.T) {
	chain := nbdns.NewHandlerChain()
	h := &answeringHandler{name: "h", ip: "10.0.0.1"}

	assert.False(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream), "empty chain")

	chain.AddHandler("example.com.", h, nbdns.PriorityUpstream)
	assert.False(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream), "non-root handler does not count")

	chain.AddHandler(".", h, nbdns.PriorityMgmtCache)
	assert.False(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream), "root handler above threshold excluded")

	chain.AddHandler(".", h, nbdns.PriorityDefault)
	assert.True(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream), "root handler at PriorityDefault included")

	chain.RemoveHandler(".", nbdns.PriorityDefault)
	assert.False(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream))

	// Primary nsgroup case: root handler lands at PriorityUpstream.
	chain.AddHandler(".", h, nbdns.PriorityUpstream)
	assert.True(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream), "root at PriorityUpstream included")
	chain.RemoveHandler(".", nbdns.PriorityUpstream)

	// Fallback case: original /etc/resolv.conf entries land at PriorityFallback.
	chain.AddHandler(".", h, nbdns.PriorityFallback)
	assert.True(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream), "root at PriorityFallback included")
	chain.RemoveHandler(".", nbdns.PriorityFallback)
	assert.False(t, chain.HasRootHandlerAtOrBelow(nbdns.PriorityUpstream))
}