package resolver

import (
	"context"
	"io"
	"net/netip"
	"testing"
	"time"

	"codeberg.org/miekg/dns"
	"codeberg.org/miekg/dns/dnsutil"
	"codeberg.org/miekg/dns/rdata"
)

func startTestServer(t *testing.T, addr string, handler dns.Handler) *dns.Server {
	t.Helper()
	ready := make(chan struct{})
	srv := &dns.Server{
		Addr:    addr,
		Net:     "udp",
		Handler: handler,
		UDPSize: 4096,
		NotifyStartedFunc: func(ctx context.Context) {
			close(ready)
		},
	}
	go func() {
		_ = srv.ListenAndServe()
	}()
	<-ready
	return srv
}

func TestLookupDirectAnswer(t *testing.T) {
	srv := startTestServer(t, "127.0.0.1:15353",
		dns.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
			resp := new(dns.Msg)
			dnsutil.SetReply(resp, req)
			resp.Authoritative = true
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.Header{
					Name:  "example.com.",
					Class: dns.ClassINET,
					TTL:   60,
				},
				A: rdata.A{Addr: netip.AddrFrom4([4]byte{127, 0, 0, 1})},
			})
			io.Copy(w, resp)
		}))
	defer srv.Shutdown(context.Background())
	time.Sleep(100 * 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.Addr != netip.AddrFrom4([4]byte{127, 0, 0, 1}) {
		t.Fatalf("expected 127.0.0.1, got %s", a.A.Addr)
	}
}

func TestLookupNXDOMAIN(t *testing.T) {
	srv := startTestServer(t, "127.0.0.1:15354",
		dns.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
			resp := new(dns.Msg)
			dnsutil.SetReply(resp, req)
			resp.Rcode = dns.RcodeNameError
			io.Copy(w, resp)
		}))
	defer srv.Shutdown(context.Background())
	time.Sleep(100 * 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.Header{Name: "example.com.", Class: dns.ClassINET, TTL: 300},
		NS:  rdata.NS{Ns: "ns1.example.com."},
	})
	msg.Extra = append(msg.Extra, &dns.A{
		Hdr: dns.Header{Name: "ns1.example.com.", Class: dns.ClassINET, TTL: 300},
		A:   rdata.A{Addr: netip.MustParseAddr("192.0.2.1")},
	})
	msg.Extra = append(msg.Extra, &dns.AAAA{
		Hdr:  dns.Header{Name: "ns1.example.com.", Class: dns.ClassINET, TTL: 300},
		AAAA: rdata.AAAA{Addr: netip.MustParseAddr("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.Header{Name: "example.com.", Class: dns.ClassINET, TTL: 300},
		NS:  rdata.NS{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")
	}
}

func TestLookupCNAME(t *testing.T) {
	callCount := 0
	srv := startTestServer(t, "127.0.0.1:15355",
		dns.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
			callCount++
			resp := new(dns.Msg)
			dnsutil.SetReply(resp, req)
			resp.Authoritative = true
			resp.Answer = append(resp.Answer,
				&dns.CNAME{
					Hdr:   dns.Header{Name: "alias.example.com.", Class: dns.ClassINET, TTL: 60},
					CNAME: rdata.CNAME{Target: "real.example.com."},
				},
			)
			if callCount > 1 {
				resp.Answer = append(resp.Answer,
					&dns.A{
						Hdr: dns.Header{Name: "real.example.com.", Class: dns.ClassINET, TTL: 60},
						A:   rdata.A{Addr: netip.AddrFrom4([4]byte{127, 0, 0, 1})},
					},
				)
			}
			io.Copy(w, resp)
		}))
	defer srv.Shutdown(context.Background())
	time.Sleep(100 * time.Millisecond)

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

	resp, err := r.Lookup(ctx, "alias.example.com.", dns.TypeA)
	if err != nil {
		t.Fatalf("CNAME lookup failed: %v", err)
	}
	if resp.Rcode != dns.RcodeSuccess {
		t.Fatalf("expected NOERROR, got %d", resp.Rcode)
	}
	hasCNAME := false
	hasA := false
	for _, rr := range resp.Answer {
		if _, ok := rr.(*dns.CNAME); ok {
			hasCNAME = true
		}
		if _, ok := rr.(*dns.A); ok {
			hasA = true
		}
	}
	if !hasCNAME {
		t.Fatal("expected CNAME record in answer")
	}
	if !hasA {
		t.Fatal("expected A record (CNAME target) in answer")
	}
}