package resolver

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

	"github.com/miekg/dns"
)

func startTestServer(t *testing.T, addr string, handler dns.Handler) *dns.Server {
	t.Helper()
	srv := &dns.Server{
		Addr:    addr,
		Net:     "udp",
		Handler: handler,
		UDPSize: 4096,
	}
	go func() {
		if err := srv.ListenAndServe(); err != nil {
		}
	}()
	return srv
}

func TestLookupDirectAnswer(t *testing.T) {
	mux := dns.NewServeMux()
	mux.HandleFunc(".", func(w dns.ResponseWriter, req *dns.Msg) {
		resp := new(dns.Msg)
		resp.SetReply(req)
		resp.Authoritative = true
		if req.Question[0].Name == "example.com." &&
			req.Question[0].Qtype == dns.TypeA {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:   "example.com.",
					Rrtype: dns.TypeA,
					Class:  dns.ClassINET,
					Ttl:    60,
				},
				A: net.IPv4(127, 0, 0, 1),
			})
		} else {
			resp.Rcode = dns.RcodeNameError
		}
		w.WriteMsg(resp)
	})

	srv := startTestServer(t, "127.0.0.1:15353", mux)
	defer srv.Shutdown()
	time.Sleep(50 * time.Millisecond)

	r := New(
		WithRootAddresses([]string{"127.0.0.1:15353"}),
		WithTimeout(500*time.Millisecond),
	)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	resp, err := r.Lookup(ctx, "example.com.", dns.TypeA)
	if err != nil {
		t.Fatalf("Lookup failed: %v", err)
	}
	if resp.Rcode != dns.RcodeSuccess {
		t.Fatalf("expected NOERROR, got %d (%s)",
			resp.Rcode, dns.RcodeToString[resp.Rcode])
	}
	if len(resp.Answer) != 1 {
		t.Fatalf("expected 1 answer, got %d", len(resp.Answer))
	}
	a, ok := resp.Answer[0].(*dns.A)
	if !ok {
		t.Fatal("expected A record")
	}
	if !a.A.Equal(net.IPv4(127, 0, 0, 1)) {
		t.Fatalf("expected 127.0.0.1, got %s", a.A)
	}
}

func TestLookupNXDOMAIN(t *testing.T) {
	mux := dns.NewServeMux()
	mux.HandleFunc(".", func(w dns.ResponseWriter, req *dns.Msg) {
		resp := new(dns.Msg)
		resp.SetReply(req)
		resp.Rcode = dns.RcodeNameError
		w.WriteMsg(resp)
	})

	srv := startTestServer(t, "127.0.0.1:15354", mux)
	defer srv.Shutdown()
	time.Sleep(50 * time.Millisecond)

	r := New(
		WithRootAddresses([]string{"127.0.0.1:15354"}),
		WithTimeout(500*time.Millisecond),
	)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	resp, err := r.Lookup(ctx, "nonexistent.xyz.", dns.TypeA)
	if err != nil {
		t.Fatalf("Lookup failed: %v", err)
	}
	if resp.Rcode != dns.RcodeNameError {
		t.Fatalf("expected NXDOMAIN, got %d", resp.Rcode)
	}
}

func TestNextServersWithGlue(t *testing.T) {
	msg := new(dns.Msg)
	msg.Authoritative = false
	msg.Ns = append(msg.Ns, &dns.NS{
		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNS, Ttl: 300},
		Ns:  "ns1.example.com.",
	})
	msg.Extra = append(msg.Extra, &dns.A{
		Hdr: dns.RR_Header{Name: "ns1.example.com.", Rrtype: dns.TypeA, Ttl: 300},
		A:   net.ParseIP("192.0.2.1").To4(),
	})
	msg.Extra = append(msg.Extra, &dns.AAAA{
		Hdr:  dns.RR_Header{Name: "ns1.example.com.", Rrtype: dns.TypeAAAA, Ttl: 300},
		AAAA: net.ParseIP("2001:db8::1"),
	})

	r := &Resolver{}
	addrs, err := r.nextServers(context.Background(), msg)
	if err != nil {
		t.Fatalf("nextServers failed: %v", err)
	}
	if len(addrs) != 1 || addrs[0] != "192.0.2.1" {
		t.Fatalf("expected [192.0.2.1], got %v", addrs)
	}
}

func TestNextServersNoGlue(t *testing.T) {
	msg := new(dns.Msg)
	msg.Authoritative = false
	msg.Ns = append(msg.Ns, &dns.NS{
		Hdr: dns.RR_Header{Name: "example.com.", Rrtype: dns.TypeNS, Ttl: 300},
		Ns:  "ns1.example.com.",
	})

	r := &Resolver{maxDelegations: 30, timeout: time.Second, retries: 1}
	_, err := r.nextServers(context.Background(), msg)
	if err == nil {
		t.Fatal("expected error when no glue and no roots")
	}
}