package resolver_test

import (
	"context"
	"ekdns/resolver"
	"net"
	"testing"
	"time"
)

// ---------------------------------------------------------------------------
// UDPQuery tests
// ---------------------------------------------------------------------------

func TestUDPQuery_Success(t *testing.T) {
	wantResponse := []byte("fake-dns-response-bytes")

	// Fake server that echoes back a fixed payload regardless of query content.
	addr := startFakeDNS(t, func(_ []byte) []byte {
		return wantResponse
	})

	ctx := context.Background()
	resp, err := resolver.UDPQuery(ctx, addr, []byte("query"), 3*time.Second)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if string(resp) != string(wantResponse) {
		t.Errorf("response: got %q, want %q", resp, wantResponse)
	}
}

func TestUDPQuery_Timeout(t *testing.T) {
	// Fake server that reads but never responds.
	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("ListenPacket: %v", err)
	}
	t.Cleanup(func() { conn.Close() })
	go func() {
		buf := make([]byte, 1232)
		for {
			_, _, err := conn.ReadFrom(buf)
			if err != nil {
				return
			}
			// Deliberately do not respond.
		}
	}()

	ctx := context.Background()
	_, err = resolver.UDPQuery(ctx, conn.LocalAddr().String(), []byte("query"), 50*time.Millisecond)
	if err == nil {
		t.Fatal("expected timeout error, got nil")
	}
}

func TestUDPQuery_ContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // Cancel immediately before making the call.

	_, err := resolver.UDPQuery(ctx, "127.0.0.1:53", []byte("query"), 3*time.Second)
	if err == nil {
		t.Fatal("expected error for cancelled context, got nil")
	}
}

func TestUDPQuery_InvalidAddress(t *testing.T) {
	ctx := context.Background()
	// Use a syntactically invalid address to provoke a dial error.
	_, err := resolver.UDPQuery(ctx, ":::invalid:::", []byte("query"), 3*time.Second)
	if err == nil {
		t.Fatal("expected error for invalid address, got nil")
	}
}