package resolver

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"
)

// UDPQuery sends a raw DNS query to addr over UDP and returns the raw response bytes.
//
// addr may be "ip" or "ip:port" — ":53" is appended if no port is specified.
// timeout sets the per-query deadline for the read/write operations.
// ctx allows early cancellation; if the context already has a deadline that
// is sooner than timeout, the context deadline takes precedence.
func UDPQuery(ctx context.Context, addr string, query []byte, timeout time.Duration) ([]byte, error) {
	if !strings.Contains(addr, ":") {
		addr = addr + ":53"
	}

	dialer := net.Dialer{Timeout: timeout}
	conn, err := dialer.DialContext(ctx, "udp", addr)
	if err != nil {
		return nil, fmt.Errorf("connecting to %s: %w", addr, err)
	}
	defer conn.Close()

	// Determine the effective deadline: use ctx deadline if sooner than timeout.
	deadline := time.Now().Add(timeout)
	if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) {
		deadline = ctxDeadline
	}
	conn.SetDeadline(deadline) //nolint:errcheck

	if _, err := conn.Write(query); err != nil {
		return nil, fmt.Errorf("sending query to %s: %w", addr, err)
	}

	buf := make([]byte, udpBufSize)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, fmt.Errorf("reading response from %s: %w", addr, err)
	}
	return buf[:n], nil
}