// Package resolver forwards queries to upstream recursive servers with
// health tracking, per-upstream circuit breakers and parallel racing.
package resolver

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"math/rand/v2"
	"net"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

// ErrNoUpstream is returned when every configured upstream is unavailable.
var ErrNoUpstream = errors.New("resolver: no healthy upstream available")

// ErrCircuitOpen is returned by a single upstream whose breaker is open.
var ErrCircuitOpen = errors.New("resolver: upstream circuit breaker open")

// lookupTimeout bounds a name resolution when building the pool, so a broken
// system resolver cannot stall a config reload.
const lookupTimeout = 2 * time.Second

// Health is an upstream's current state.
type Health string

const (
	// HealthUp means the last probes and queries succeeded.
	HealthUp Health = "up"
	// HealthDegraded means recent failures but still serving.
	HealthDegraded Health = "degraded"
	// HealthDown means failing or the breaker is open.
	HealthDown Health = "down"
	// HealthUnknown means not probed yet.
	HealthUnknown Health = "unknown"
)

// upstreamStats is everything mutable about an upstream. It lives behind a
// pointer shared by successive Upstream values for the same server, so a
// reload keeps the live health state while the immutable Upstream itself is
// replaced under the readers' feet.
type upstreamStats struct {
	queries     atomic.Uint64
	failures    atomic.Uint64
	consecutive atomic.Uint32
	latencySum  atomic.Uint64
	latencyN    atomic.Uint64

	// ring keeps the last 512 latencies in ms for the UI's p95.
	mu   sync.Mutex
	ring []float64

	breakerOpenUntil atomic.Int64
	health           atomic.Pointer[Health]
	lastErr          atomic.Pointer[string]
}

func newUpstreamStats() *upstreamStats {
	s := &upstreamStats{}
	s.health.Store(ptr(HealthUnknown))
	return s
}

// Upstream is one configured server: its address, its dial-time
// configuration and a pointer to the live stats it shares across reloads.
//
// An Upstream is immutable after construction. Rebind builds a fresh Upstream
// for every server and swaps the whole pool state with one atomic store, so a
// query in flight keeps using a consistent, unchanging view.
type Upstream struct {
	cfg config.UpstreamServer
	// addr is the address actually dialled: a bootstrap IP for TLS (SNI keeps
	// the hostname) and the literal address otherwise.
	addr string
	// resolved holds the bootstrap addresses behind a named upstream.
	resolved []string
	stats    *upstreamStats
}

// Addr is the address the resolver dials.
func (u *Upstream) Addr() string { return u.addr }

// Config returns the upstream's configuration entry.
func (u *Upstream) Config() config.UpstreamServer { return u.cfg }

// Resolved returns the addresses a named upstream resolved to, if any.
func (u *Upstream) Resolved() []string { return u.resolved }

// Health returns the current health state.
func (u *Upstream) Health() Health {
	if u == nil || u.stats == nil {
		return HealthUnknown
	}
	if h := u.stats.health.Load(); h != nil {
		return *h
	}
	return HealthUnknown
}

func (u *Upstream) setHealth(h Health) { u.stats.health.Store(&h) }

// SetError records the last error message for the UI.
func (u *Upstream) SetError(err error) {
	if err == nil {
		u.stats.lastErr.Store(nil)
		return
	}
	s := err.Error()
	u.stats.lastErr.Store(&s)
}

// LastError returns the last error message, or "".
func (u *Upstream) LastError() string {
	if p := u.stats.lastErr.Load(); p != nil {
		return *p
	}
	return ""
}

// Success records a successful query.
func (u *Upstream) Success(d time.Duration) {
	s := u.stats
	s.queries.Add(1)
	s.consecutive.Store(0)
	s.latencySum.Add(uint64(d.Microseconds()))
	s.latencyN.Add(1)
	u.recordLatency(d)
	u.SetError(nil)
	// A working upstream closes its own breaker. Without this the breaker
	// could only ever be rearmed by the cooldown expiring, so a server that
	// recovered in a millisecond stayed excluded from every pick() for the
	// whole cooldown window.
	s.breakerOpenUntil.Store(0)
	if u.Health() != HealthUp {
		u.setHealth(HealthUp)
	}
}

// Failure records a failed query and trips the breaker once the threshold is
// crossed.
func (u *Upstream) Failure(err error, threshold int, cooldown time.Duration) {
	s := u.stats
	s.queries.Add(1)
	s.failures.Add(1)
	n := s.consecutive.Add(1)
	u.SetError(err)
	if threshold > 0 && int(n) >= threshold {
		s.breakerOpenUntil.Store(time.Now().Add(cooldown).UnixNano())
		u.setHealth(HealthDown)
		return
	}
	if u.Health() != HealthDown {
		u.setHealth(HealthDegraded)
	}
}

// BreakerOpen reports whether the upstream's breaker is open.
func (u *Upstream) BreakerOpen() bool {
	return u.stats.breakerOpenUntil.Load() > time.Now().UnixNano()
}

func (u *Upstream) recordLatency(d time.Duration) {
	ms := float64(d.Microseconds()) / 1000
	s := u.stats
	s.mu.Lock()
	if len(s.ring) < 512 {
		s.ring = append(s.ring, ms)
	} else {
		copy(s.ring, s.ring[1:])
		s.ring[511] = ms
	}
	s.mu.Unlock()
}

// LatencyStats returns average and p95 latency in milliseconds.
func (u *Upstream) LatencyStats() (avg, p95 float64) {
	s := u.stats
	n := s.latencyN.Load()
	if n == 0 {
		return 0, 0
	}
	avg = float64(s.latencySum.Load()) / float64(n) / 1000
	s.mu.Lock()
	samples := append([]float64(nil), s.ring...)
	s.mu.Unlock()
	if len(samples) == 0 {
		return avg, 0
	}
	sort.Float64s(samples)
	idx := int(float64(len(samples)) * 0.95)
	if idx >= len(samples) {
		idx = len(samples) - 1
	}
	return avg, samples[idx]
}

// Stats is the JSON-friendly upstream state for the API.
type Stats struct {
	Addr        string  `json:"addr"`
	Proto       string  `json:"proto"`
	ServerName  string  `json:"server_name"`
	Weight      int     `json:"weight"`
	Comment     string  `json:"comment,omitempty"`
	Health      Health  `json:"health"`
	LatencyMS   float64 `json:"latency_ms"`
	AvgMS       float64 `json:"avg_ms"`
	P95MS       float64 `json:"p95_ms"`
	FailRate    float64 `json:"fail_rate"`
	Queries     uint64  `json:"queries"`
	Failures    uint64  `json:"failures"`
	BreakerOpen bool    `json:"breaker_open"`
	LastError   string  `json:"last_error,omitempty"`
}

// Stats snapshots one upstream.
func (u *Upstream) Stats() Stats {
	if u == nil {
		return Stats{}
	}
	q := u.stats.queries.Load()
	f := u.stats.failures.Load()
	var rate float64
	if q > 0 {
		rate = float64(f) / float64(q)
	}
	avg, p95 := u.LatencyStats()
	return Stats{
		Addr:        u.addr,
		Proto:       u.cfg.Proto,
		ServerName:  u.cfg.ServerName,
		Weight:      u.cfg.Weight,
		Comment:     u.cfg.Comment,
		Health:      u.Health(),
		LatencyMS:   p95,
		AvgMS:       avg,
		P95MS:       p95,
		FailRate:    rate,
		Queries:     q,
		Failures:    f,
		BreakerOpen: u.BreakerOpen(),
		LastError:   u.LastError(),
	}
}

// poolState is an immutable snapshot of the pool: the config it was built
// from, the upstreams, and the address index used for per-client upstream
// overrides. Readers take the pointer once and are then immune to reloads.
type poolState struct {
	cfg      config.Upstream
	ups      []*Upstream
	defaults []*Upstream
	routes   map[string][]*Upstream
	// byAddr indexes the upstreams by every address they answer to, so an ACL
	// group naming either the configured Addr or the resolved one works.
	byAddr map[string][]*Upstream
}

func (st *poolState) all() []*Upstream {
	if st == nil {
		return nil
	}
	return st.ups
}

// Pool is the set of upstreams plus the policy that selects among them.
type Pool struct {
	state atomic.Pointer[poolState]
	// clients are memoised per upstream: one UDP socket, one TCP dialer, one
	// TLS session pool. Reusing them is what keeps the per-query cost low; a
	// dns.Client is safe for concurrent use.
	clients sync.Map // *Upstream -> *dns.Client

	// mu serialises Rebind so two overlapping reloads cannot interleave their
	// health hand-over. It is never taken on the query path.
	mu sync.Mutex
	// rng backs the "random" policy.
	rngMu sync.Mutex
	rng   *rand.Rand

	// udpSize is the read buffer for UDP clients. It must be at least as large
	// as the EDNS0 buffer we advertise upstream: a UDP datagram bigger than
	// this buffer is read truncated, fails to unpack, and turns into a
	// SERVFAIL. See clientFor.
	udpSize uint16
}

// DefaultClientUDPSize is the read buffer used when the caller does not say
// otherwise. It matches the historical hard-coded client buffer.
const DefaultClientUDPSize uint16 = 1232

// NewPool builds a pool from config. Bootstraps are used for named upstreams;
// anything unresolved falls back to the system resolver.
func NewPool(cfg config.Upstream) *Pool {
	return NewPoolWithUDPSize(cfg, DefaultClientUDPSize)
}

// NewPoolWithUDPSize builds a pool whose UDP clients read up to udpSize bytes.
// Callers that advertise an EDNS0 buffer to their upstreams must pass that
// same value here, or every answer larger than it is lost.
func NewPoolWithUDPSize(cfg config.Upstream, udpSize uint16) *Pool {
	// 512 is the floor a resolver may assume; below it we would truncate
	// replies a plain client asked for.
	if udpSize < dns.MinMsgSize {
		udpSize = dns.MinMsgSize
	}
	p := &Pool{
		rng:     rand.New(rand.NewPCG(uint64(time.Now().UnixNano()), 0x9e3779b9)),
		udpSize: udpSize,
	}
	p.state.Store(buildState(cfg, nil, false))
	return p
}

// buildState turns an upstream config into an immutable pool snapshot. When
// carry is non-nil, servers that keep the same address and protocol also keep
// their live health, latency and breaker state.
func buildState(cfg config.Upstream, carry *poolState, lookup bool) *poolState {
	bootstraps := make(map[string][]string, len(cfg.Bootstraps))
	for _, b := range cfg.Bootstraps {
		bootstraps[b.Name] = append([]string(nil), b.Addresses...)
	}
	prev := make(map[string]*Upstream)
	if carry != nil {
		for _, u := range carry.ups {
			prev[u.cfg.Addr+"/"+u.cfg.Proto] = u
		}
	}

	st := &poolState{cfg: cfg, ups: make([]*Upstream, 0, len(cfg.Servers))}
	for _, s := range cfg.Servers {
		addr, resolved := resolveAddr(s, bootstraps, lookup)
		up := &Upstream{cfg: s, addr: addr, resolved: resolved}
		if old, ok := prev[s.Addr+"/"+s.Proto]; ok {
			up.stats = old.stats
		} else {
			up.stats = newUpstreamStats()
		}
		st.ups = append(st.ups, up)
	}
	st.defaults = append([]*Upstream(nil), st.ups...)
	st.routes = make(map[string][]*Upstream)
	for _, z := range cfg.ForwardZones {
		zone := strings.ToLower(dns.Fqdn(z.Zone))
		for _, s := range z.Servers {
			addr, resolved := resolveAddr(s, bootstraps, lookup)
			up := &Upstream{cfg: s, addr: addr, resolved: resolved}
			if old := prev[s.Addr+"/"+s.Proto]; old != nil {
				up.stats = old.stats
			} else {
				up.stats = newUpstreamStats()
			}
			st.ups = append(st.ups, up)
			st.routes[zone] = append(st.routes[zone], up)
		}
	}
	st.byAddr = make(map[string][]*Upstream, len(st.ups)*2)
	for _, up := range st.ups {
		for _, a := range append([]string{up.cfg.Addr, up.addr}, up.resolved...) {
			if a == "" {
				continue
			}
			st.byAddr[a] = append(st.byAddr[a], up)
		}
	}
	return st
}

// resolveAddr returns the address to dial for a server plus the bootstrap
// addresses behind it, if the configured host is a name.
func resolveAddr(s config.UpstreamServer, bootstraps map[string][]string, lookup bool) (addr string, resolved []string) {
	host, port, err := net.SplitHostPort(s.Addr)
	if err != nil {
		// Not host:port — hand it to the dialer verbatim.
		return s.Addr, nil
	}
	if net.ParseIP(host) != nil {
		return s.Addr, nil
	}
	if addrs := bootstraps[host]; len(addrs) > 0 {
		return net.JoinHostPort(addrs[0], port), addrs
	}
	if lookup {
		ctx, cancel := context.WithTimeout(context.Background(), lookupTimeout)
		defer cancel()
		addrs, err := net.DefaultResolver.LookupHost(ctx, host)
		if err == nil && len(addrs) > 0 {
			return net.JoinHostPort(addrs[0], port), addrs
		}
	}
	return s.Addr, nil
}

// Config returns the pool's upstream config.
func (p *Pool) Config() config.Upstream {
	if st := p.state.Load(); st != nil {
		return st.cfg
	}
	return config.Upstream{}
}

// All returns every upstream.
func (p *Pool) All() []*Upstream { return p.state.Load().all() }

// Defaults returns the servers eligible for queries without a forwarding route.
func (p *Pool) Defaults() []*Upstream { return p.state.Load().defaults }

// Get returns an upstream by configured or resolved address, or nil.
func (p *Pool) Get(addr string) *Upstream {
	for _, u := range p.All() {
		if u.addr == addr || u.cfg.Addr == addr {
			return u
		}
	}
	return nil
}

// Healthy returns upstreams whose breaker is closed, in configured order.
func (p *Pool) Healthy() []*Upstream {
	ups := p.All()
	out := make([]*Upstream, 0, len(ups))
	for _, u := range ups {
		if !u.BreakerOpen() {
			out = append(out, u)
		}
	}
	return out
}

// Stats snapshots every upstream.
func (p *Pool) Stats() []Stats {
	ups := p.All()
	out := make([]Stats, 0, len(ups))
	for _, u := range ups {
		out = append(out, u.Stats())
	}
	return out
}

// clientFor returns the memoised dns.Client for an upstream.
func (p *Pool) clientFor(u *Upstream) *dns.Client {
	if c, ok := p.clients.Load(u); ok {
		return c.(*dns.Client)
	}
	nc := &dns.Client{Net: "udp", UDPSize: p.udpSize, Timeout: 3 * time.Second}
	switch u.cfg.Proto {
	case "tls":
		nc = &dns.Client{
			Net: "tcp-tls",
			TLSConfig: &tls.Config{
				ServerName:         u.cfg.ServerName,
				InsecureSkipVerify: u.cfg.InsecureSkipVerify,
				MinVersion:         tls.VersionTLS12,
				NextProtos:         []string{"dot"},
			},
			Dialer:  &net.Dialer{Timeout: 3 * time.Second},
			Timeout: 5 * time.Second,
		}
	case "tcp":
		nc = &dns.Client{Net: "tcp", Timeout: 5 * time.Second}
	}
	actual, _ := p.clients.LoadOrStore(u, nc)
	return actual.(*dns.Client)
}

// Exchange forwards q to the pool and returns the first acceptable answer.
//
// want restricts the query to the upstreams an ACL group named, by configured
// or resolved address; an empty list means the whole pool. timeout bounds one
// attempt.
func (p *Pool) Exchange(ctx context.Context, q *dns.Msg, want []string, timeout time.Duration) (*dns.Msg, string, error) {
	st := p.state.Load()
	candidates := st.defaults
	// Walk complete labels, most specific first. A matched route never falls
	// back to public resolvers, including when every target is unavailable.
	matched := false
	if len(q.Question) > 0 {
		name := strings.ToLower(dns.Fqdn(q.Question[0].Name))
		for name != "" {
			if route, ok := st.routes[name]; ok {
				candidates = route
				matched = true
				break
			}
			dot := strings.IndexByte(name, '.')
			if dot < 0 {
				break
			}
			name = name[dot+1:]
		}
	}
	if !matched && len(want) > 0 {
		candidates = p.restrict(st, want)
	}
	if len(candidates) == 0 {
		return nil, "", ErrNoUpstream
	}
	chosen := p.pick(st, candidates, st.cfg.MaxAttempts)
	if len(chosen) == 0 {
		return nil, "", ErrNoUpstream
	}
	if len(chosen) == 1 {
		return p.exchangeOne(ctx, st.cfg, chosen[0], q, timeout)
	}
	if st.cfg.Policy != "parallel" && st.cfg.Policy != "" {
		// Failover and random are *selection* policies: pick() has already
		// ordered (or shuffled) the candidates, so the only correct thing to
		// do with the list is try it in order. Fanning out here would send
		// every failover query to the secondary as well, which is the one
		// thing failover exists to avoid.
		var lastErr error
		for _, up := range chosen {
			msg, addr, err := p.exchangeOne(ctx, st.cfg, up, q, timeout)
			if err == nil && msg != nil && msg.Rcode != dns.RcodeServerFailure {
				return msg, addr, nil
			}
			if err != nil {
				lastErr = err
			} else {
				lastErr = fmt.Errorf("upstream %s returned SERVFAIL", addr)
			}
		}
		if lastErr == nil {
			lastErr = ErrNoUpstream
		}
		return nil, "", lastErr
	}

	// Parallel: send to all at once and take the first real answer. The
	// context is cancelled on return so we do not leak a request to a slow
	// upstream after another one has answered.
	type res struct {
		msg *dns.Msg
		up  string
		err error
	}
	results := make(chan res, len(chosen))
	sub, cancel := context.WithCancel(ctx)
	defer cancel()
	for _, up := range chosen {
		go func(up *Upstream) {
			msg, addr, err := p.exchangeOne(sub, st.cfg, up, q, timeout)
			results <- res{msg: msg, up: addr, err: err}
		}(up)
	}
	var lastErr error
	for i := 0; i < len(chosen); i++ {
		select {
		case r := <-results:
			switch {
			case r.err != nil:
				lastErr = r.err
			case r.msg != nil && r.msg.Rcode != dns.RcodeServerFailure:
				return r.msg, r.up, nil
			default:
				lastErr = fmt.Errorf("upstream %s returned SERVFAIL", r.up)
			}
		case <-ctx.Done():
			return nil, "", ctx.Err()
		}
	}
	if lastErr == nil {
		lastErr = ErrNoUpstream
	}
	return nil, "", lastErr
}

// restrict resolves an ACL group's upstream list to upstreams, keeping the
// configured order and dropping names that match nothing. When nothing in the
// list resolves the full pool is used, so a typo cannot black-hole a client.
func (p *Pool) restrict(st *poolState, want []string) []*Upstream {
	seen := make(map[*Upstream]bool, len(want))
	out := make([]*Upstream, 0, len(want))
	for _, name := range want {
		for _, up := range st.byAddr[name] {
			if !seen[up] {
				seen[up] = true
				out = append(out, up)
			}
		}
	}
	if len(out) == 0 {
		return st.all()
	}
	return out
}

// pick filters out open breakers (keeping everything if that empties the
// list) and applies the pool policy.
func (p *Pool) pick(st *poolState, ups []*Upstream, maxAttempts int) []*Upstream {
	open := make([]*Upstream, 0, len(ups))
	for _, u := range ups {
		if !u.BreakerOpen() {
			open = append(open, u)
		}
	}
	if len(open) == 0 {
		open = ups
	}
	if maxAttempts <= 0 || maxAttempts > len(open) {
		maxAttempts = len(open)
	}
	switch st.cfg.Policy {
	case "failover":
		return open[:maxAttempts]
	case "random":
		p.rngMu.Lock()
		defer p.rngMu.Unlock()
		p.rng.Shuffle(len(open), func(i, j int) { open[i], open[j] = open[j], open[i] })
		return open[:maxAttempts]
	default:
		return open
	}
}

// exchangeOne sends q to a single upstream, retrying over TCP when the UDP
// answer comes back truncated (RFC 1035 §4.2.1).
func (p *Pool) exchangeOne(ctx context.Context, cfg config.Upstream, up *Upstream, q *dns.Msg, timeout time.Duration) (*dns.Msg, string, error) {
	if timeout <= 0 {
		timeout = 2 * time.Second
	}
	client := p.clientFor(up)
	resp, rtt, err := client.ExchangeContext(ctx, q.Copy(), up.addr)
	if err != nil {
		up.Failure(err, thresholdFor(up, cfg.BreakerThreshold), cfg.BreakerCooldown.D())
		return nil, up.addr, err
	}
	up.Success(rtt)
	if !resp.Truncated {
		return resp, up.addr, nil
	}
	// Retry over TCP for the full answer. An upstream that is UDP-only just
	// fails here and we return the truncated answer; the client retries over
	// TCP itself.
	tcpClient := &dns.Client{Net: "tcp", Timeout: timeout}
	resp2, rtt2, err2 := tcpClient.ExchangeContext(ctx, q.Copy(), up.addr)
	if err2 != nil || resp2 == nil || resp2.Truncated {
		if err2 == nil {
			err2 = errors.New("still truncated over TCP")
		}
		up.Failure(fmt.Errorf("truncated answer, TCP retry failed: %w", err2),
			thresholdFor(up, cfg.BreakerThreshold), cfg.BreakerCooldown.D())
		return resp, up.addr, nil
	}
	up.Success(rtt2)
	return resp2, up.addr, nil
}

// HealthCheck probes every upstream with a cheap query, refreshing health
// state even when no client traffic is flowing. It returns the number of
// upstreams that answered.
func (p *Pool) HealthCheck(ctx context.Context, timeout time.Duration) int {
	st := p.state.Load()
	probe := new(dns.Msg)
	probe.SetQuestion(".", dns.TypeNS)
	probe.RecursionDesired = true
	ok := 0
	for _, up := range st.all() {
		if up.BreakerOpen() {
			continue
		}
		c := p.clientFor(up)
		_, rtt, err := c.ExchangeContext(ctx, probe, up.addr)
		if err != nil {
			up.Failure(err, thresholdFor(up, st.cfg.BreakerThreshold), st.cfg.BreakerCooldown.D())
			continue
		}
		up.Success(rtt)
		ok++
	}
	return ok
}

func thresholdFor(up *Upstream, def int) int {
	if up.cfg.MaxFails > 0 {
		return up.cfg.MaxFails
	}
	return def
}

// Rebind applies a new upstream configuration, keeping the live health state
// of upstreams whose address and protocol are unchanged.
//
// It builds a whole new immutable state and swaps it in with a single atomic
// store: queries already in flight finish against the state they started
// with, and the next query sees the new one. There is no window in which half
// the upstreams are old and half new.
func (p *Pool) Rebind(cfg config.Upstream) {
	p.mu.Lock()
	defer p.mu.Unlock()
	old := p.state.Load()
	next := buildState(cfg, old, true)
	p.state.Store(next)
	// Drop the memoised clients of upstreams that are gone; their shells are
	// only reachable from in-flight queries, which let them go on their own.
	live := make(map[*Upstream]bool, len(next.all()))
	for _, u := range next.all() {
		live[u] = true
	}
	for _, u := range old.all() {
		if !live[u] {
			p.clients.Delete(u)
		}
	}
}

// String renders the pool for logs.
func (p *Pool) String() string {
	ups := p.All()
	addrs := make([]string, 0, len(ups))
	for _, u := range ups {
		addrs = append(addrs, u.addr)
	}
	return "pool[" + strings.Join(addrs, " ") + "]"
}

func ptr[T any](v T) *T { return &v }