// 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")

// 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"
)

// Upstream is one configured server plus its live health.
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 current bootstrap addresses when the configured
	// address is a name.
	resolved []string

	queries     atomic.Uint64
	failures    atomic.Uint64
	consecutive atomic.Uint32
	latencySum  atomic.Uint64 // microseconds
	latencyN    atomic.Uint64
	p95         []float64 // recent latencies in ms, for the UI
	p95mu       sync.Mutex
	p95ring     []float64

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

// 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 }

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

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

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

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

// Success records a successful query.
func (u *Upstream) Success(d time.Duration) {
	u.queries.Add(1)
	u.consecutive.Store(0)
	u.latencySum.Add(uint64(d.Microseconds()))
	u.latencyN.Add(1)
	u.recordLatency(d)
	u.SetError(nil)
	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) {
	u.queries.Add(1)
	u.failures.Add(1)
	n := u.consecutive.Add(1)
	u.SetError(err)
	if threshold > 0 && int(n) >= threshold {
		u.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 {
	until := u.breakerOpenUntil.Load()
	return until > time.Now().UnixNano()
}

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

// LatencyStats returns average and p95 latency in milliseconds.
func (u *Upstream) LatencyStats() (avg, p95 float64) {
	n := u.latencyN.Load()
	if n == 0 {
		return 0, 0
	}
	avg = float64(u.latencySum.Load()) / float64(n) / 1000
	u.p95mu.Lock()
	samples := append([]float64(nil), u.p95ring...)
	u.p95mu.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 {
	q := u.queries.Load()
	f := u.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(),
	}
}

// Pool is the set of upstreams plus the policy that selects among them.
type Pool struct {
	cfg   config.Upstream
	mu    sync.RWMutex
	ups   []*Upstream
	byTag map[string][]*Upstream
	// clients are reused per upstream: one UDP socket, one TCP dialer, one TLS
	// session pool. Reusing them is what keeps the per-query cost low.
	clients map[*Upstream]*dns.Client
	// bootstrap resolves names to addresses.
	bootstrap map[string][]string

	rng *rand.Rand
	mu2 sync.Mutex
}

// NewPool builds a pool from config. bootstrapIPs may be empty; then upstream
// hostnames must be literal IPs (or resolved by the OS).
func NewPool(cfg config.Upstream) *Pool {
	p := &Pool{
		cfg:       cfg,
		byTag:     make(map[string][]*Upstream),
		clients:   make(map[*Upstream]*dns.Client),
		bootstrap: make(map[string][]string),
		rng:       rand.New(rand.NewPCG(uint64(time.Now().UnixNano()), 0x9e3779b9)),
	}
	for _, b := range cfg.Bootstraps {
		p.bootstrap[b.Name] = append([]string(nil), b.Addresses...)
	}
	for _, s := range cfg.Servers {
		up := &Upstream{cfg: s}
		up.setHealth(HealthUnknown)
		p.ups = append(p.ups, up)
	}
	p.resolveAll()
	return p
}

func (p *Pool) resolveAll() {
	for _, up := range p.ups {
		host, port, err := net.SplitHostPort(up.cfg.Addr)
		if err != nil {
			continue
		}
		if net.ParseIP(host) != nil {
			up.addr = up.cfg.Addr
			continue
		}
		if addrs := p.bootstrap[host]; len(addrs) > 0 {
			up.resolved = addrs
			up.addr = net.JoinHostPort(addrs[0], port)
			continue
		}
		// Fall back to the system resolver.
		ips, err := net.LookupIP(host)
		if err != nil || len(ips) == 0 {
			up.addr = up.cfg.Addr
			continue
		}
		strs := make([]string, 0, len(ips))
		for _, ip := range ips {
			strs = append(strs, ip.String())
		}
		up.resolved = strs
		up.addr = net.JoinHostPort(strs[0], port)
	}
}

// Config returns the pool's upstream config.
func (p *Pool) Config() config.Upstream { return p.cfg }

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

// Get returns an upstream by address, or nil.
func (p *Pool) Get(addr string) *Upstream {
	for _, u := range p.ups {
		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 {
	out := make([]*Upstream, 0, len(p.ups))
	for _, u := range p.ups {
		if !u.BreakerOpen() {
			out = append(out, u)
		}
	}
	return out
}

// Stats snapshots every upstream.
func (p *Pool) Stats() []Stats {
	out := make([]Stats, 0, len(p.ups))
	for _, u := range p.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 {
	p.mu.RLock()
	c, ok := p.clients[u]
	p.mu.RUnlock()
	if ok {
		return c
	}
	var nc *dns.Client
	switch u.cfg.Proto {
	case "tls":
		tlsCfg := &tls.Config{
			ServerName:         u.cfg.ServerName,
			InsecureSkipVerify: u.cfg.InsecureSkipVerify,
			MinVersion:         tls.VersionTLS12,
			NextProtos:         []string{"dot"},
		}
		nc = &dns.Client{
			Net:       "tcp-tls",
			TLSConfig: tlsCfg,
			Dialer:    &net.Dialer{Timeout: 3 * time.Second},
			Timeout:   5 * time.Second,
		}
	case "tcp":
		nc = &dns.Client{Net: "tcp", Timeout: 5 * time.Second}
	default:
		nc = &dns.Client{Net: "udp", UDPSize: 1232, Timeout: 3 * time.Second}
	}
	p.mu.Lock()
	p.clients[u] = nc
	p.mu.Unlock()
	return nc
}

// selectFor returns the upstreams to try, honouring the policy. tags, if
// non-empty, restricts the pool to a group's own upstreams.
func (p *Pool) selectFor(policy string, maxAttempts int) []*Upstream {
	healthy := p.Healthy()
	if len(healthy) == 0 {
		// Every breaker is open: try them all anyway rather than failing the
		// query outright. A blip should not become an outage.
		healthy = append(healthy, p.ups...)
	}
	if len(healthy) == 0 {
		return nil
	}
	if maxAttempts <= 0 || maxAttempts > len(healthy) {
		maxAttempts = len(healthy)
	}
	switch policy {
	case "failover":
		// Configured order, and rotate past the unhealthy ones.
		return healthy[:maxAttempts]
	case "random":
		p.mu2.Lock()
		p.rng.Shuffle(len(healthy), func(i, j int) { healthy[i], healthy[j] = healthy[j], healthy[i] })
		p.mu2.Unlock()
		return healthy[:maxAttempts]
	default: // parallel
		return healthy
	}
}

// Exchange forwards q to the pool and returns the first acceptable answer.
//
// policy is the pool's configured policy; tags restricts the pool to a
// per-client upstream override. edns0Size is the client's advertised buffer,
// used to decide TCP fallback on truncation.
func (p *Pool) Exchange(ctx context.Context, q *dns.Msg, tags []string, timeout time.Duration) (*dns.Msg, string, error) {
	candidates := p.ups
	if len(tags) > 0 {
		if tagged := p.byTag[strings.Join(tags, ",")]; len(tagged) > 0 {
			candidates = tagged
		}
	}
	if len(candidates) == 0 {
		return nil, "", ErrNoUpstream
	}
	attempts := p.cfg.MaxAttempts
	chosen := selectFrom(candidates, p.cfg.Policy, attempts)
	if len(chosen) == 0 {
		return nil, "", ErrNoUpstream
	}

	if len(chosen) == 1 {
		msg, up, err := p.exchangeOne(ctx, chosen[0], q, timeout)
		return msg, up, err
	}

	// Parallel: send to all at once, take the first answer that is not a
	// transport failure. Cancellable so we do not leak requests 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, 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:
			if r.err == nil && r.msg != nil && !isTransportFailure(r.msg.Rcode) {
				return r.msg, r.up, nil
			}
			if r.err != nil {
				lastErr = r.err
			} else if r.msg != nil {
				lastErr = fmt.Errorf("upstream %s returned %s", r.up, dns.RcodeToString[r.msg.Rcode])
			}
		case <-ctx.Done():
			return nil, "", ctx.Err()
		}
	}
	if lastErr == nil {
		lastErr = ErrNoUpstream
	}
	return nil, "", lastErr
}

func selectFrom(ups []*Upstream, policy string, maxAttempts int) []*Upstream {
	// Filter out open breakers but keep everything if that empties the list.
	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 policy {
	case "failover":
		return open[:maxAttempts]
	case "random":
		rand.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, up *Upstream, q *dns.Msg, timeout time.Duration) (*dns.Msg, string, error) {
	if timeout <= 0 {
		timeout = 2 * time.Second
	}
	client := p.clientFor(up)
	msg := q.Copy()
	start := time.Now()
	resp, rtt, err := client.ExchangeContext(ctx, msg, up.addr)
	up.Success(rtt)
	if err != nil {
		return nil, up.addr, err
	}
	_ = start
	if !resp.Truncated {
		return resp, up.addr, nil
	}
	// Retry over TCP to get the full answer. Some upstreams are configured
	// UDP-only in which case this simply fails and we use the TC answer.
	tcpClient := &dns.Client{Net: "tcp", Timeout: timeout}
	resp2, rtt2, err2 := tcpClient.ExchangeContext(ctx, q.Copy(), up.addr)
	if err2 == nil && resp2 != nil && !resp2.Truncated {
		up.Success(rtt2)
		return resp2, up.addr, nil
	}
	up.Failure(fmt.Errorf("truncated answer, TCP retry failed: %w", err2), up.breakerThreshold(), p.cfg.BreakerCooldown.D())
	// Return the truncated answer; the client will retry over TCP itself.
	return resp, up.addr, nil
}

func isTransportFailure(rcode int) bool {
	return rcode == dns.RcodeServerFailure
}

// breakerThreshold is the effective failure threshold for an upstream.
func (u *Upstream) breakerThreshold() int {
	if u.cfg.MaxFails > 0 {
		return u.cfg.MaxFails
	}
	return 0
}

// 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 {
	probe := new(dns.Msg)
	probe.SetQuestion(".", dns.TypeNS)
	probe.RecursionDesired = true
	ok := 0
	for _, up := range p.ups {
		if up.BreakerOpen() {
			continue
		}
		c := p.clientFor(up)
		_, rtt, err := c.ExchangeContext(ctx, probe, up.addr)
		if err != nil {
			up.Failure(err, p.breakerThresholdFor(up), p.cfg.BreakerCooldown.D())
			continue
		}
		up.Success(rtt)
		ok++
	}
	return ok
}

func (p *Pool) breakerThresholdFor(u *Upstream) int {
	if t := u.breakerThreshold(); t > 0 {
		return t
	}
	return p.cfg.BreakerThreshold
}

// Rebind applies a new upstream configuration in place, keeping the live
// health state of upstreams whose address is unchanged.
func (p *Pool) Rebind(cfg config.Upstream) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.cfg = cfg
	old := make(map[string]*Upstream, len(p.ups))
	for _, u := range p.ups {
		old[u.cfg.Addr+"/"+u.cfg.Proto] = u
	}
	next := make([]*Upstream, 0, len(cfg.Servers))
	for _, s := range cfg.Servers {
		if u, ok := old[s.Addr+"/"+s.Proto]; ok {
			u.cfg = s
			next = append(next, u)
			delete(old, s.Addr+"/"+s.Proto)
			continue
		}
		u := &Upstream{cfg: s}
		u.setHealth(HealthUnknown)
		next = append(next, u)
	}
	p.ups = next
	// Drop clients for removed upstreams.
	for u := range p.clients {
		if old[u.cfg.Addr+"/"+u.cfg.Proto] != nil {
			delete(p.clients, u)
		}
	}
	p.resolveAllLocked()
}

func (p *Pool) resolveAllLocked() {
	for _, b := range p.cfg.Bootstraps {
		p.bootstrap[b.Name] = append([]string(nil), b.Addresses...)
	}
	for _, up := range p.ups {
		host, port, err := net.SplitHostPort(up.cfg.Addr)
		if err != nil {
			continue
		}
		if net.ParseIP(host) != nil {
			up.addr = up.cfg.Addr
			continue
		}
		if addrs := p.bootstrap[host]; len(addrs) > 0 {
			up.resolved = addrs
			up.addr = net.JoinHostPort(addrs[0], port)
			continue
		}
		up.addr = up.cfg.Addr
	}
}