package e2e

import (
	"net"
	"testing"
	"time"
)

// MockDNSServer is a minimal UDP server that responds to any DNS query
// with a hardcoded A record for example.com.
type MockDNSServer struct {
	conn *net.UDPConn
}

func StartMockDNSServer(t *testing.T) (*MockDNSServer, int) {
	addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("failed to resolve UDP address: %v", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		t.Fatalf("failed to start mock DNS server: %v", err)
	}

	return &MockDNSServer{conn: conn}, conn.LocalAddr().(*net.UDPAddr).Port
}

func (s *MockDNSServer) Close() {
	_ = s.conn.Close()
}

// ListenAndRespond runs in a loop and responds to the first valid-looking DNS query.
// It returns the response as a channel to allow the test to synchronize.
func (s *MockDNSServer) ListenAndRespond(timeout time.Duration) (chan bool, error) {
	respChan := make(chan bool, 1)

	go func() {
		buf := make([]byte, 1024)
		_ = s.conn.SetReadDeadline(time.Now().Add(timeout))
		n, remoteAddr, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			respChan <- false
			return
		}

		if n < 4 {
			respChan <- false
			return
		}

		// Echo the packet back to the client as a basic "I heard you" confirmation.
		_, _ = s.conn.WriteToUDP(buf[:n], remoteAddr)
		respChan <- true
	}()

	return respChan, nil
}