package resolver

import (
	"context"
	"errors"
	"net"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

// fakeOpt configures a fakeUpstream.
type fakeOpt struct {
	ip          string        // A record handed out in the answer
	delay       time.Duration // how long the handler pretends to think
	truncateUDP bool          // answer TC=1 over UDP, forcing the client onto TCP
}

// fakeUpstream is a real dns.Server on 127.0.0.1, listening for UDP and TCP on
// the same port, so tests exercise the pool's real client path end to end.
type fakeUpstream struct {
	addr string
	opt  fakeOpt

	udpHits atomic.Int64
	tcpHits atomic.Int64
	// broken makes the handler answer with a message that cannot be unpacked,
	// which is the fastest way to make the client fail a real exchange.
	broken atomic.Bool
}

// newFakeUpstream starts a fake upstream and stops it when the test ends.
func newFakeUpstream(t *testing.T, opt fakeOpt) *fakeUpstream {
	t.Helper()
	if opt.ip == "" {
		opt.ip = "203.0.113.1"
	}
	lc, pc := listenSamePort(t)
	f := &fakeUpstream{addr: lc.Addr().String(), opt: opt}

	// ActivateAndServe serves either the PacketConn or the Listener, never
	// both, so the fake runs one server per transport.
	srvs := []*dns.Server{
		{PacketConn: pc, Handler: dns.HandlerFunc(f.handle)},
		{Listener: lc, Handler: dns.HandlerFunc(f.handle)},
	}
	stop := make(chan struct{})
	var wg sync.WaitGroup
	for _, srv := range srvs {
		wg.Add(1)
		go func(srv *dns.Server) {
			defer wg.Done()
			_ = srv.ActivateAndServe()
		}(srv)
	}
	go func() {
		<-stop
		for _, srv := range srvs {
			_ = srv.Shutdown()
		}
	}()
	t.Cleanup(func() {
		// close(stop) must happen before wg.Wait(): the serve goroutines only
		// returns once the watcher has shut the server down, so waiting first
		// would hang the whole package instead of failing this test.
		close(stop)
		wg.Wait()
	})
	return f
}

// listenSamePort binds a loopback port for TCP and mirrors it for UDP, which is
// what a real recursive server does and what the truncation retry needs.
func listenSamePort(t *testing.T) (net.Listener, net.PacketConn) {
	t.Helper()
	for attempt := 0; attempt < 5; attempt++ {
		lc, err := net.Listen("tcp", "127.0.0.1:0")
		if err != nil {
			t.Fatalf("listen tcp: %v", err)
		}
		_, port, err := net.SplitHostPort(lc.Addr().String())
		if err != nil {
			_ = lc.Close()
			t.Fatalf("split %s: %v", lc.Addr(), err)
		}
		pc, err := net.ListenPacket("udp", "127.0.0.1:"+port)
		if err != nil {
			// Someone grabbed the UDP port in between; try a fresh one.
			_ = lc.Close()
			continue
		}
		return lc, pc
	}
	t.Fatal("no loopback port was free for UDP and TCP at once")
	return nil, nil
}

func (f *fakeUpstream) handle(w dns.ResponseWriter, q *dns.Msg) {
	tcp := w.RemoteAddr().Network() == "tcp"
	if tcp {
		f.tcpHits.Add(1)
	} else {
		f.udpHits.Add(1)
	}
	if f.broken.Load() {
		// A header claiming five questions with no question in the body: the
		// client cannot unpack it and reports an error, which is exactly the
		// failure path a broken upstream takes in production. The query id is
		// echoed, otherwise the client drops the reply and waits for its
		// timeout instead of reporting the error.
		bad := []byte{byte(q.Id >> 8), byte(q.Id), 0x80, 0, 0, 1, 0, 0, 0, 0, 0, 0}
		_, _ = w.Write(append(bad, 3, 'a', 'b'))
		return
	}
	if d := f.opt.delay; d > 0 {
		time.Sleep(d)
	}
	m := new(dns.Msg)
	m.SetReply(q)
	m.Truncated = !tcp && f.opt.truncateUDP
	if q.Question[0].Qtype == dns.TypeA && !m.Truncated {
		m.Answer = []dns.RR{&dns.A{
			Hdr: dns.RR_Header{Name: q.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
			A:   net.ParseIP(f.opt.ip).To4(),
		}}
	}
	_ = w.WriteMsg(m)
}

// closedPort returns a loopback port with nothing behind it, so an exchange to
// it fails immediately instead of waiting for a timeout.
func closedPort(t *testing.T) string {
	t.Helper()
	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := pc.LocalAddr().String()
	if err := pc.Close(); err != nil {
		t.Fatalf("close: %v", err)
	}
	return addr
}

func udpServer(addr string) config.UpstreamServer {
	return config.UpstreamServer{Addr: addr, Proto: "udp", Weight: 1}
}

// poolCfg is a parallel pool over the given servers.
func poolCfg(servers ...config.UpstreamServer) config.Upstream {
	return config.Upstream{
		Servers:          servers,
		Policy:           "parallel",
		MaxAttempts:      len(servers),
		BreakerThreshold: 3,
		BreakerCooldown:  config.Duration(30 * time.Second),
	}
}

func query(name string) *dns.Msg {
	q := new(dns.Msg)
	q.SetQuestion(dns.Fqdn(name), dns.TypeA)
	q.RecursionDesired = true
	return q
}

func testCtx(t *testing.T) context.Context {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	t.Cleanup(cancel)
	return ctx
}

// ask runs one exchange and insists it succeeded.
func ask(t *testing.T, p *Pool, want []string, name string) (*dns.Msg, string) {
	t.Helper()
	resp, addr, err := p.Exchange(testCtx(t), query(name), want, time.Second)
	if err != nil {
		t.Fatalf("exchange %s: %v", name, err)
	}
	return resp, addr
}

// answerIP returns the first A record of a reply.
func answerIP(t *testing.T, m *dns.Msg) string {
	t.Helper()
	for _, rr := range m.Answer {
		if a, ok := rr.(*dns.A); ok {
			return a.A.String()
		}
	}
	t.Fatalf("no A record in reply %v", m)
	return ""
}

func TestCircuitBreakerOpensAfterConsecutiveFailures(t *testing.T) {
	dead := closedPort(t)
	cfg := config.Upstream{
		Servers:          []config.UpstreamServer{udpServer(dead)},
		Policy:           "parallel",
		MaxAttempts:      1,
		BreakerThreshold: 2,
		BreakerCooldown:  config.Duration(30 * time.Second),
	}
	p := NewPool(cfg)
	up := p.Get(dead)

	// First failure: still serving, just degraded.
	if _, _, err := p.Exchange(testCtx(t), query("a.test."), nil, time.Second); err == nil {
		t.Fatal("exchange to a dead upstream unexpectedly succeeded")
	}
	if up.BreakerOpen() {
		t.Error("breaker open after one failure with a threshold of 2")
	}
	if up.Health() != HealthDegraded {
		t.Errorf("health = %q, want %q after one failure", up.Health(), HealthDegraded)
	}

	// Second consecutive failure trips it.
	if _, _, err := p.Exchange(testCtx(t), query("b.test."), nil, time.Second); err == nil {
		t.Fatal("second exchange to a dead upstream unexpectedly succeeded")
	}
	if !up.BreakerOpen() {
		t.Fatalf("breaker closed after %d consecutive failures with a threshold of 2", up.Stats().Failures)
	}
	if up.Health() != HealthDown {
		t.Errorf("health = %q, want %q while the breaker is open", up.Health(), HealthDown)
	}
	if got := p.Healthy(); len(got) != 0 {
		t.Errorf("Healthy() = %v, want no upstreams while the breaker is open", p)
	}
	s := p.Stats()[0]
	if !s.BreakerOpen || s.Failures != 2 || s.Queries != 2 || s.FailRate != 1 {
		t.Errorf("stats = %+v, want 2 queries, 2 failures, fail rate 1 and an open breaker", s)
	}
	if s.LastError == "" {
		t.Error("stats carry no last error after two failures")
	}
}

func TestCircuitBreakerClosesAfterCooldown(t *testing.T) {
	dead := closedPort(t)
	const cooldown = 60 * time.Millisecond
	cfg := config.Upstream{
		Servers:          []config.UpstreamServer{udpServer(dead)},
		Policy:           "parallel",
		MaxAttempts:      1,
		BreakerThreshold: 1,
		BreakerCooldown:  config.Duration(cooldown),
	}
	p := NewPool(cfg)
	up := p.Get(dead)

	if _, _, err := p.Exchange(testCtx(t), query("a.test."), nil, time.Second); err == nil {
		t.Fatal("exchange to a dead upstream unexpectedly succeeded")
	}
	if !up.BreakerOpen() {
		t.Fatal("breaker stayed closed with a threshold of 1")
	}

	// Still open well inside the cooldown window.
	time.Sleep(cooldown / 3)
	if !up.BreakerOpen() {
		t.Fatal("breaker closed before the cooldown elapsed")
	}

	// The package exposes no clock seam, so the cooldown boundary can only be
	// crossed with a real sleep; cooldown is the shortest duration the config
	// accepts and this keeps the test far below a second.
	time.Sleep(cooldown)
	if up.BreakerOpen() {
		t.Fatal("breaker still open after the cooldown expired")
	}
	if got := p.Healthy(); len(got) != 1 || got[0] != up {
		t.Errorf("Healthy() = %v, want the upstream back in the list", p)
	}
	if up.Health() != HealthDown {
		t.Errorf("health = %q, want %q until the next successful probe", up.Health(), HealthDown)
	}
}

// TestFailedExchangeDoesNotMarkUpstreamHealthy is the regression test for a
// bug where Success(rtt) ran before the error was checked, so a failed exchange
// pinned the upstream at HealthUp and reset the consecutive-failure counter.
func TestFailedExchangeDoesNotMarkUpstreamHealthy(t *testing.T) {
	f := newFakeUpstream(t, fakeOpt{ip: "203.0.113.5"})
	cfg := poolCfg(udpServer(f.addr))
	cfg.BreakerThreshold = 2
	p := NewPool(cfg)
	up := p.Get(f.addr)

	// A good exchange first, so "up" is a state the upstream can lose.
	ask(t, p, nil, "ok.test.")
	if up.Health() != HealthUp {
		t.Fatalf("health = %q after a good exchange, want %q", up.Health(), HealthUp)
	}

	f.broken.Store(true)
	if _, _, err := p.Exchange(testCtx(t), query("bad.test."), nil, time.Second); err == nil {
		t.Fatal("an unparsable answer must surface as an error")
	}
	if got := up.Health(); got != HealthDegraded {
		t.Errorf("health = %q after a failed exchange, want %q (never %q)", got, HealthDegraded, HealthUp)
	}
	if s := up.Stats(); s.Queries != 2 || s.Failures != 1 || s.FailRate != 0.5 {
		t.Errorf("stats = %+v, want 2 queries, 1 failure, fail rate 0.5", s)
	}
	if up.LastError() == "" {
		t.Error("a failed exchange recorded no error")
	}
	if up.BreakerOpen() {
		t.Error("a single failure tripped a threshold of 2")
	}

	// The failure counts: one more consecutive failure opens the breaker even
	// though an earlier success reset the counter.
	if _, _, err := p.Exchange(testCtx(t), query("bad2.test."), nil, time.Second); err == nil {
		t.Fatal("an unparsable answer must surface as an error")
	}
	if !up.BreakerOpen() {
		t.Error("two consecutive failures did not open the breaker")
	}

	// And the upstream recovers as soon as it answers again.
	f.broken.Store(false)
	p.Rebind(poolCfg(udpServer(f.addr)))
	ask(t, p, nil, "ok2.test.")
	if up.Health() != HealthUp {
		t.Errorf("health = %q after recovery, want %q", up.Health(), HealthUp)
	}
	if up.BreakerOpen() {
		t.Error("breaker still open after a successful exchange")
	}
	if up.LastError() != "" {
		t.Errorf("last error %q survived a success", up.LastError())
	}
}

func TestParallelPolicyAnswersFromTheFastestUpstream(t *testing.T) {
	fast := newFakeUpstream(t, fakeOpt{ip: "203.0.113.10"})
	slow1 := newFakeUpstream(t, fakeOpt{ip: "203.0.113.11", delay: 250 * time.Millisecond})
	slow2 := newFakeUpstream(t, fakeOpt{ip: "203.0.113.12", delay: 400 * time.Millisecond})
	p := NewPool(poolCfg(udpServer(fast.addr), udpServer(slow1.addr), udpServer(slow2.addr)))

	start := time.Now()
	resp, addr := ask(t, p, nil, "race.test.")
	elapsed := time.Since(start)

	if addr != fast.addr {
		t.Errorf("pool answered from %s, want the fastest upstream %s", addr, fast.addr)
	}
	if ip := answerIP(t, resp); ip != "203.0.113.10" {
		t.Errorf("answer = %s, want the fast upstream's record", ip)
	}
	if elapsed >= 250*time.Millisecond {
		t.Errorf("exchange took %s, so it waited for a slow upstream instead of racing", elapsed)
	}
	// The fan-out dispatches the slow upstreams in the background, so whether
	// they were reached before the winner cancelled the context is a genuine
	// race — asserting a hit count made this flaky. What matters is that the
	// answer came from the fast upstream and that nothing waited on a slow
	// one, both asserted above.
	if hits := slow1.udpHits.Load() + slow2.udpHits.Load(); hits > 2 {
		t.Errorf("slow upstreams were queried %d times for a single exchange, want at most one each", hits)
	}
}

// TestParallelPolicyConcurrentExchanges is a race-detector stress test: a burst
// of queries sharing the pool's memoised dns.Client and mutable health state.
func TestParallelPolicyConcurrentExchanges(t *testing.T) {
	a := newFakeUpstream(t, fakeOpt{ip: "203.0.113.20"})
	b := newFakeUpstream(t, fakeOpt{ip: "203.0.113.21"})
	p := NewPool(poolCfg(udpServer(a.addr), udpServer(b.addr)))

	// One client per upstream, reused for every query.
	up := p.Get(a.addr)
	if p.clientFor(up) != p.clientFor(up) {
		t.Error("clientFor built a second dns.Client for the same upstream")
	}

	const goroutines, perGoroutine = 16, 4
	var wg sync.WaitGroup
	errs := make(chan error, goroutines*perGoroutine)
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func(g int) {
			defer wg.Done()
			for i := 0; i < perGoroutine; i++ {
				resp, addr, err := p.Exchange(testCtx(t), query("burst.test."), nil, time.Second)
				if err != nil {
					errs <- err
					return
				}
				if addr != a.addr && addr != b.addr {
					errs <- errors.New("answered by " + addr)
					return
				}
				if ip := answerIP(t, resp); ip != "203.0.113.20" && ip != "203.0.113.21" {
					errs <- errors.New("unexpected answer " + ip)
					return
				}
			}
		}(g)
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		t.Error(err)
	}

	// The parallel policy asks *both* upstreams and keeps the first answer, so
	// every query produces two requests. The upper bound holds because the
	// winner cancels the context the loser is on, which can drop a request
	// that was already in flight — hence "at least one, at most two".
	queries := goroutines * perGoroutine
	if got := a.udpHits.Load() + b.udpHits.Load(); got <= int64(queries) || got > 2*int64(queries) {
		t.Errorf("upstreams answered %d queries, want between %d and %d", got, queries, 2*queries)
	}
	if s := p.Get(a.addr).Stats(); s.Queries == 0 {
		t.Error("no query was recorded against the first upstream")
	}
}

func TestTruncatedUDPAnswerIsRetriedOverTCP(t *testing.T) {
	f := newFakeUpstream(t, fakeOpt{ip: "203.0.113.30", truncateUDP: true})
	p := NewPool(poolCfg(udpServer(f.addr)))

	resp, addr := ask(t, p, nil, "big.test.")
	if addr != f.addr {
		t.Errorf("answered by %s, want %s", addr, f.addr)
	}
	if resp.Truncated {
		t.Fatal("pool handed back the truncated UDP answer instead of retrying over TCP")
	}
	if ip := answerIP(t, resp); ip != "203.0.113.30" {
		t.Errorf("answer = %s, want the full record from the TCP retry", ip)
	}
	if f.udpHits.Load() == 0 {
		t.Error("the query never went out over UDP")
	}
	if f.tcpHits.Load() == 0 {
		t.Error("the truncated answer was not retried over TCP")
	}
}

func TestFailoverPolicyWalksTheListInOrder(t *testing.T) {
	first := newFakeUpstream(t, fakeOpt{ip: "203.0.113.40"})
	second := newFakeUpstream(t, fakeOpt{ip: "203.0.113.41"})
	cfg := config.Upstream{
		Servers:          []config.UpstreamServer{udpServer(first.addr), udpServer(second.addr)},
		Policy:           "failover",
		MaxAttempts:      2,
		BreakerThreshold: 3,
		BreakerCooldown:  config.Duration(30 * time.Second),
	}
	p := NewPool(cfg)

	// The primary serves every query while it is healthy.
	for i := 0; i < 3; i++ {
		resp, addr := ask(t, p, nil, "failover.test.")
		if addr != first.addr {
			t.Fatalf("query %d answered by %s, want the primary %s", i, addr, first.addr)
		}
		if ip := answerIP(t, resp); ip != "203.0.113.40" {
			t.Fatalf("query %d answered %s, want the primary's record", i, ip)
		}
	}
	if second.udpHits.Load() != 0 {
		t.Errorf("the secondary took %d queries while the primary was healthy", second.udpHits.Load())
	}

	// Once the primary starts failing the secondary takes over.
	first.broken.Store(true)
	resp, addr := ask(t, p, nil, "failover.test.")
	if addr != second.addr {
		t.Errorf("answered by %s, want the secondary %s once the primary fails", addr, second.addr)
	}
	if ip := answerIP(t, resp); ip != "203.0.113.41" {
		t.Errorf("answer = %s, want the secondary's record", ip)
	}
	if s := p.Get(first.addr).Stats(); s.Failures != 1 {
		t.Errorf("primary stats = %+v, want exactly one recorded failure", s)
	}

	// And the pool goes back to the primary when it recovers.
	first.broken.Store(false)
	if _, addr := ask(t, p, nil, "failover.test."); addr != first.addr {
		t.Errorf("answered by %s after the primary recovered, want %s", addr, first.addr)
	}
}

// TestRebindUnderConcurrentQueries hammers Rebind while queries, health
// snapshots and latency reads run in parallel, so -race covers the copy-on-write
// hand-over. Afterwards the live state of an unchanged server must have
// survived and a changed server list must be visible.
func TestRebindUnderConcurrentQueries(t *testing.T) {
	a := newFakeUpstream(t, fakeOpt{ip: "203.0.113.50"})
	b := newFakeUpstream(t, fakeOpt{ip: "203.0.113.51"})
	c := newFakeUpstream(t, fakeOpt{ip: "203.0.113.52"})

	only := poolCfg(udpServer(a.addr))
	base := poolCfg(udpServer(a.addr), udpServer(b.addr))
	alt := poolCfg(udpServer(a.addr), udpServer(c.addr))

	p := NewPool(only)
	ask(t, p, nil, "rebind.test.")
	before := p.All()[0]
	if before.Health() != HealthUp {
		t.Fatalf("health = %q before the storm, want %q", before.Health(), HealthUp)
	}
	queriesBefore := before.Stats().Queries

	stop := make(chan struct{})
	var seen sync.Map // address -> true, every list the readers ever observed
	var wg sync.WaitGroup
	for r := 0; r < 4; r++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-stop:
					return
				default:
				}
				for _, u := range p.All() {
					seen.Store(u.Addr(), true)
				}
				p.Healthy()
				p.Stats()
				if _, _, err := p.Exchange(testCtx(t), query("storm.test."), nil, time.Second); err != nil {
					t.Errorf("exchange during rebind: %v", err)
					return
				}
				runtime.Gosched()
			}
		}()
	}
	for i := 0; i < 25; i++ {
		p.Rebind(alt)
		runtime.Gosched()
		p.Rebind(base)
	}
	close(stop)
	wg.Wait()

	// The unchanged server kept its live health and latency state. Rebind
	// always builds a fresh immutable Upstream shell — the *stats* are what
	// carry over, so that is what is asserted here, not the pointer.
	p.Rebind(only)
	after := p.All()[0]
	if after.stats != before.stats {
		t.Error("Rebind did not hand the live stats to the new upstream value")
	}
	if got := after.Health(); got != HealthUp {
		t.Errorf("health = %q after the storm, want %q", got, HealthUp)
	}
	if got := after.Stats().Queries; got < queriesBefore {
		t.Errorf("query counter went backwards across Rebind: %d < %d", got, queriesBefore)
	}
	if avg, p95 := after.LatencyStats(); avg <= 0 || p95 <= 0 {
		t.Errorf("latency stats = avg %v p95 %v, want the samples recorded before the storm", avg, p95)
	}

	// The changed list is visible.
	p.Rebind(alt)
	if p.Get(c.addr) == nil {
		t.Error("the added upstream is missing after Rebind")
	}
	if p.Get(b.addr) != nil {
		t.Error("the removed upstream is still in the pool")
	}
	if got := len(p.All()); got != 2 {
		t.Errorf("pool holds %d upstreams, want 2", got)
	}
	for _, addr := range []string{b.addr, c.addr} {
		if _, ok := seen.Load(addr); !ok {
			t.Errorf("no reader ever observed %s while Rebind swapped the list", addr)
		}
	}
	if !strings.Contains(p.String(), c.addr) {
		t.Errorf("pool %s does not mention the new upstream %s", p, c.addr)
	}
}

func TestExchangeRestrictsToNamedUpstreams(t *testing.T) {
	a := newFakeUpstream(t, fakeOpt{ip: "203.0.113.60"})
	b := newFakeUpstream(t, fakeOpt{ip: "203.0.113.61"})
	p := NewPool(poolCfg(udpServer(a.addr), udpServer(b.addr)))

	_, addr := ask(t, p, []string{b.addr}, "restricted.test.")
	if addr != b.addr {
		t.Errorf("answered by %s, want the named upstream %s", addr, b.addr)
	}
	if a.udpHits.Load() != 0 {
		t.Errorf("the unnamed upstream took %d queries", a.udpHits.Load())
	}

	// A name that matches nothing falls back to the whole pool rather than
	// black-holing the client.
	_, addr = ask(t, p, []string{"nope.invalid"}, "fallback.test.")
	if addr != a.addr && addr != b.addr {
		t.Errorf("answered by %s, want any configured upstream", addr)
	}

	// An empty pool has nothing to offer.
	empty := NewPool(config.Upstream{Policy: "parallel"})
	if _, _, err := empty.Exchange(testCtx(t), query("nowhere.test."), nil, time.Second); !errors.Is(err, ErrNoUpstream) {
		t.Errorf("err = %v, want %v", err, ErrNoUpstream)
	}
}

func TestHealthCheckProbesEveryUpstream(t *testing.T) {
	live := newFakeUpstream(t, fakeOpt{ip: "203.0.113.70"})
	dead := closedPort(t)
	cfg := poolCfg(udpServer(live.addr), udpServer(dead))
	cfg.BreakerThreshold = 1
	p := NewPool(cfg)

	if n := p.HealthCheck(testCtx(t), time.Second); n != 1 {
		t.Errorf("HealthCheck reported %d healthy upstreams, want 1", n)
	}
	if got := p.Get(live.addr).Health(); got != HealthUp {
		t.Errorf("health of the live upstream = %q, want %q", got, HealthUp)
	}
	if !p.Get(dead).BreakerOpen() {
		t.Error("the dead upstream was not taken out of rotation")
	}
}