//go:build darwin || linux

package daemon

import (
	"fmt"
	"net"
	"os"
	"sync"
	"time"
)

const (
	// maxDNSPacketSize is the maximum UDP packet size we accept.
	// DNS can theoretically be up to 65535 bytes, but practically much smaller.
	maxDNSPacketSize = 4096

	// upstreamTimeout is the time we wait for a response from the upstream DNS server.
	upstreamTimeout = 5 * time.Second
)

// DNSRelay is a UDP DNS relay that forwards DNS queries from sandboxed processes
// to a configured upstream DNS server. It operates as a simple packet relay without
// parsing DNS protocol contents.
type DNSRelay struct {
	udpConn    *net.UDPConn
	targetAddr string // upstream DNS server address (host:port)
	listenAddr string // address we're listening on
	wg         sync.WaitGroup
	done       chan struct{}
	debug      bool
}

// NewDNSRelay creates a new DNS relay that listens on listenAddr and forwards
// queries to dnsAddr. The listenAddr will typically be "127.0.0.2:53" (loopback alias).
// The dnsAddr must be in "host:port" format (e.g. "1.1.1.1:53").
func NewDNSRelay(listenAddr, dnsAddr string, debug bool) (*DNSRelay, error) {
	// Validate the upstream DNS address is parseable as host:port.
	targetHost, targetPort, err := net.SplitHostPort(dnsAddr)
	if err != nil {
		return nil, fmt.Errorf("invalid DNS address %q: %w", dnsAddr, err)
	}
	if targetHost == "" {
		return nil, fmt.Errorf("invalid DNS address %q: empty host", dnsAddr)
	}
	if targetPort == "" {
		return nil, fmt.Errorf("invalid DNS address %q: empty port", dnsAddr)
	}

	// Resolve and bind the listen address.
	udpAddr, err := net.ResolveUDPAddr("udp", listenAddr)
	if err != nil {
		return nil, fmt.Errorf("failed to resolve listen address %q: %w", listenAddr, err)
	}

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return nil, fmt.Errorf("failed to bind UDP socket on %q: %w", listenAddr, err)
	}

	return &DNSRelay{
		udpConn:    conn,
		targetAddr: dnsAddr,
		listenAddr: conn.LocalAddr().String(),
		done:       make(chan struct{}),
		debug:      debug,
	}, nil
}

// ListenAddr returns the actual address the relay is listening on.
// This is useful when port 0 was used to get an ephemeral port.
func (d *DNSRelay) ListenAddr() string {
	return d.listenAddr
}

// Start begins the DNS relay loop. It reads incoming UDP packets from the
// listening socket and spawns a goroutine per query to forward it to the
// upstream DNS server and relay the response back.
func (d *DNSRelay) Start() error {
	if d.debug {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Listening on %s, forwarding to %s\n", d.listenAddr, d.targetAddr)
	}

	d.wg.Add(1)
	go d.readLoop()

	return nil
}

// Stop shuts down the DNS relay. It signals the read loop to stop, closes the
// listening socket, and waits for all in-flight queries to complete.
func (d *DNSRelay) Stop() {
	close(d.done)
	_ = d.udpConn.Close()
	d.wg.Wait()

	if d.debug {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Stopped\n")
	}
}

// readLoop is the main loop that reads incoming DNS queries from the listening socket.
func (d *DNSRelay) readLoop() {
	defer d.wg.Done()

	buf := make([]byte, maxDNSPacketSize)
	for {
		n, clientAddr, err := d.udpConn.ReadFromUDP(buf)
		if err != nil {
			select {
			case <-d.done:
				// Shutting down, expected error from closed socket.
				return
			default:
				fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Read error: %v\n", err)
				continue
			}
		}

		if n == 0 {
			continue
		}

		// Copy the packet data so the buffer can be reused immediately.
		query := make([]byte, n)
		copy(query, buf[:n])

		d.wg.Add(1)
		go d.handleQuery(query, clientAddr)
	}
}

// handleQuery forwards a single DNS query to the upstream server and relays
// the response back to the original client. It creates a dedicated UDP connection
// to the upstream server to avoid multiplexing complexity.
func (d *DNSRelay) handleQuery(query []byte, clientAddr *net.UDPAddr) {
	defer d.wg.Done()

	if d.debug {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Query from %s (%d bytes)\n", clientAddr, len(query))
	}

	// Create a dedicated UDP connection to the upstream DNS server.
	upstreamConn, err := net.Dial("udp", d.targetAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Failed to connect to upstream %s: %v\n", d.targetAddr, err)
		return
	}
	defer upstreamConn.Close() //nolint:errcheck // best-effort cleanup of per-query UDP connection

	// Send the query to the upstream server.
	if _, err := upstreamConn.Write(query); err != nil {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Failed to send query to upstream: %v\n", err)
		return
	}

	// Wait for the response with a timeout.
	if err := upstreamConn.SetReadDeadline(time.Now().Add(upstreamTimeout)); err != nil {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Failed to set read deadline: %v\n", err)
		return
	}

	resp := make([]byte, maxDNSPacketSize)
	n, err := upstreamConn.Read(resp)
	if err != nil {
		if d.debug {
			fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Upstream response error: %v\n", err)
		}
		return
	}

	// Send the response back to the original client.
	if _, err := d.udpConn.WriteToUDP(resp[:n], clientAddr); err != nil {
		// The listening socket may have been closed during shutdown.
		select {
		case <-d.done:
			return
		default:
			fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Failed to send response to %s: %v\n", clientAddr, err)
		}
	}

	if d.debug {
		fmt.Fprintf(os.Stderr, "[greywall:dns-relay] Response to %s (%d bytes)\n", clientAddr, n)
	}
}