package dns

import (
	"context"
	"errors"
	"net"
	"strings"
	"sync/atomic"
	"time"
)

// ServerPool load-balances DNS lookups across multiple resolvers using
// round-robin, and fails over to the next server on a timeout or server error.
// Definitive answers (NXDOMAIN) are returned immediately and are not retried on
// every server.
// defaultAttemptTimeout bounds each individual server attempt so a dead
// resolver costs a couple of seconds before failover, not the entire lookup
// budget.
const defaultAttemptTimeout = 2 * time.Second

type ServerPool struct {
	resolvers      []Resolver
	next           atomic.Uint64
	attemptTimeout time.Duration
}

func (p *ServerPool) withAttemptTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
	t := p.attemptTimeout
	if t <= 0 {
		t = defaultAttemptTimeout
	}
	return context.WithTimeout(ctx, t)
}

// NewServerPool builds a pool from "host", "host:port" or "system" entries.
// An empty list uses the system resolver.
func NewServerPool(servers []string) *ServerPool {
	var resolvers []Resolver
	for _, s := range servers {
		s = strings.TrimSpace(s)
		switch {
		case s == "":
			continue
		case s == "system" || s == "local":
			return &ServerPool{resolvers: []Resolver{&GoResolver{R: net.DefaultResolver}}}
		default:
			if !strings.Contains(s, ":") {
				s += ":53"
			}
			resolvers = append(resolvers, &GoResolver{R: newServerResolver(s)})
		}
	}
	if len(resolvers) == 0 {
		resolvers = append(resolvers, &GoResolver{R: net.DefaultResolver})
	}
	return &ServerPool{resolvers: resolvers}
}

func newServerResolver(server string) *net.Resolver {
	return &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
			if network != "tcp" {
				network = "udp"
			}
			d := net.Dialer{Timeout: 5 * time.Second}
			return d.DialContext(ctx, network, server)
		},
	}
}

func (p *ServerPool) pick() int {
	return int((p.next.Add(1) - 1) % uint64(len(p.resolvers)))
}

// definitive reports whether err is a real DNS answer (NXDOMAIN) that should not
// be retried against other servers.
func definitive(err error) bool {
	var dnsErr *net.DNSError
	return errors.As(err, &dnsErr) && dnsErr.IsNotFound
}

func (p *ServerPool) LookupCNAME(ctx context.Context, host string) (string, error) {
	n := len(p.resolvers)
	start := p.pick()
	var last error
	for i := 0; i < n; i++ {
		actx, cancel := p.withAttemptTimeout(ctx)
		cname, err := p.resolvers[(start+i)%n].LookupCNAME(actx, host)
		cancel()
		if err == nil {
			return cname, nil
		}
		if definitive(err) {
			return "", err
		}
		last = err
	}
	return "", last
}

func (p *ServerPool) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	n := len(p.resolvers)
	start := p.pick()
	var last error
	for i := 0; i < n; i++ {
		actx, cancel := p.withAttemptTimeout(ctx)
		ips, err := p.resolvers[(start+i)%n].LookupIP(actx, host)
		cancel()
		if err == nil {
			return ips, nil
		}
		if definitive(err) {
			return nil, err
		}
		last = err
	}
	return nil, last
}