// Package resolver performs recursive resolution by forwarding to configured
// upstream servers.
//
// Forwarding rather than full iteration is a deliberate choice for an
// appliance of this size: it is far simpler to get right, it inherits the
// upstream's own cache and DNSSEC validation, and it avoids shipping a root
// hints file that goes stale. The interface is narrow enough that a full
// iterative resolver could be dropped in behind it later.
package resolver

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"

	"github.com/owen/vibedns/internal/config"
)

// Config controls upstream behaviour. It is replaced wholesale on change.
type Config struct {
	Upstreams     []string
	Timeout       time.Duration
	Retries       int
	Strategy      string
	DNSSEC        bool
	EDNSUDPSize   uint16
	MaxConcurrent int
}

// Upstream tracks one configured server and its observed health.
type Upstream struct {
	Addr  string
	order int

	// latencyUS is an exponentially weighted moving average in microseconds.
	latencyUS atomic.Int64
	queries   atomic.Int64
	failures  atomic.Int64
	// consecutive failures; after failureThreshold the server is rested.
	consecutive atomic.Int64
	downUntil   atomic.Int64 // unix nanos
	lastError   atomic.Value // string
	lastUsed    atomic.Int64 // unix nanos
}

const (
	failureThreshold = 3
	restPeriod       = 20 * time.Second
	// initialLatency seeds the EWMA so an unqueried server is neither
	// unfairly preferred nor permanently ignored.
	initialLatencyUS = 50_000
)

func newUpstream(addr string, order int) *Upstream {
	u := &Upstream{Addr: addr, order: order}
	u.latencyUS.Store(initialLatencyUS)
	u.lastError.Store("")
	return u
}

func (u *Upstream) healthy() bool {
	until := u.downUntil.Load()
	return until == 0 || time.Now().UnixNano() >= until
}

func (u *Upstream) recordSuccess(d time.Duration) {
	// EWMA with alpha = 1/4, cheap and stable enough for server selection.
	prev := u.latencyUS.Load()
	next := (prev*3 + d.Microseconds()) / 4
	u.latencyUS.Store(next)
	u.queries.Add(1)
	u.consecutive.Store(0)
	u.downUntil.Store(0)
	u.lastUsed.Store(time.Now().UnixNano())
}

func (u *Upstream) recordFailure(err error) {
	u.failures.Add(1)
	u.queries.Add(1)
	u.lastUsed.Store(time.Now().UnixNano())
	if err != nil {
		u.lastError.Store(err.Error())
	}
	if u.consecutive.Add(1) >= failureThreshold {
		u.downUntil.Store(time.Now().Add(restPeriod).UnixNano())
	}
}

// Status is a point-in-time view of one upstream for the resolver page.
type Status struct {
	Address   string     `json:"address"`
	Healthy   bool       `json:"healthy"`
	Queries   int64      `json:"queries"`
	Failures  int64      `json:"failures"`
	LatencyMS float64    `json:"latency_ms"`
	LastError string     `json:"last_error,omitempty"`
	LastUsed  *time.Time `json:"last_used,omitempty"`
}

// Status renders the upstream's current state.
func (u *Upstream) Status() Status {
	s := Status{
		Address:   u.Addr,
		Healthy:   u.healthy(),
		Queries:   u.queries.Load(),
		Failures:  u.failures.Load(),
		LatencyMS: float64(u.latencyUS.Load()) / 1000,
	}
	if v, ok := u.lastError.Load().(string); ok {
		s.LastError = v
	}
	if n := u.lastUsed.Load(); n > 0 {
		t := time.Unix(0, n)
		s.LastUsed = &t
	}
	return s
}

// Resolver forwards queries to upstream servers.
type Resolver struct {
	mu        sync.RWMutex
	cfg       Config
	upstreams []*Upstream

	rrCounter atomic.Uint64
	sem       chan struct{}
	semMu     sync.Mutex

	udpClient *dns.Client
	tcpClient *dns.Client

	queries   atomic.Int64
	failures  atomic.Int64
	truncated atomic.Int64
	latencyUS atomic.Int64 // EWMA across all upstreams
}

// Common resolver errors surfaced to the operator.
var (
	ErrNoUpstreams = errors.New("no upstream resolvers are configured")
	ErrAllFailed   = errors.New("every upstream resolver failed to answer")
)

// New creates a resolver with the given configuration.
func New(cfg Config) *Resolver {
	r := &Resolver{}
	r.udpClient = &dns.Client{Net: "udp"}
	r.tcpClient = &dns.Client{Net: "tcp"}
	r.SetConfig(cfg)
	return r
}

// SetConfig replaces the resolver configuration. Upstreams that are still
// present keep their health statistics so a settings change does not discard
// what we have learned about them.
func (r *Resolver) SetConfig(cfg Config) {
	if cfg.Timeout <= 0 {
		cfg.Timeout = 2 * time.Second
	}
	if cfg.MaxConcurrent <= 0 {
		cfg.MaxConcurrent = 256
	}
	if cfg.EDNSUDPSize == 0 {
		cfg.EDNSUDPSize = 1232
	}

	r.mu.Lock()
	prev := map[string]*Upstream{}
	for _, u := range r.upstreams {
		prev[u.Addr] = u
	}
	ups := make([]*Upstream, 0, len(cfg.Upstreams))
	for i, addr := range cfg.Upstreams {
		if u, ok := prev[addr]; ok {
			u.order = i
			ups = append(ups, u)
			continue
		}
		ups = append(ups, newUpstream(addr, i))
	}
	r.cfg = cfg
	r.upstreams = ups
	r.udpClient.Timeout = cfg.Timeout
	r.tcpClient.Timeout = cfg.Timeout
	r.mu.Unlock()

	r.semMu.Lock()
	r.sem = make(chan struct{}, cfg.MaxConcurrent)
	r.semMu.Unlock()
}

// Config returns the active configuration.
func (r *Resolver) Config() Config {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.cfg
}

// Upstreams returns the configured upstreams in configuration order.
func (r *Resolver) Upstreams() []*Upstream {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]*Upstream, len(r.upstreams))
	copy(out, r.upstreams)
	sort.Slice(out, func(i, j int) bool { return out[i].order < out[j].order })
	return out
}

// Statuses renders every upstream's health for the UI.
func (r *Resolver) Statuses() []Status {
	ups := r.Upstreams()
	out := make([]Status, 0, len(ups))
	for _, u := range ups {
		out = append(out, u.Status())
	}
	return out
}

// order returns the upstreams to try, in the order dictated by the strategy.
func (r *Resolver) order() []*Upstream {
	r.mu.RLock()
	strategy := r.cfg.Strategy
	ups := make([]*Upstream, len(r.upstreams))
	copy(ups, r.upstreams)
	r.mu.RUnlock()

	switch strategy {
	case config.StrategyRandom:
		rand.Shuffle(len(ups), func(i, j int) { ups[i], ups[j] = ups[j], ups[i] })
	case config.StrategyRoundRobin:
		if len(ups) > 1 {
			n := int(r.rrCounter.Add(1)-1) % len(ups)
			ups = append(ups[n:], ups[:n]...)
		}
	case config.StrategyFastest:
		sort.SliceStable(ups, func(i, j int) bool {
			return ups[i].latencyUS.Load() < ups[j].latencyUS.Load()
		})
	default: // sequential
		sort.SliceStable(ups, func(i, j int) bool { return ups[i].order < ups[j].order })
	}

	// Regardless of strategy, servers that are resting go last rather than
	// being removed: if every server is unhealthy we must still try something.
	sort.SliceStable(ups, func(i, j int) bool {
		return ups[i].healthy() && !ups[j].healthy()
	})
	return ups
}

// Result carries a forwarded answer and where it came from.
type Result struct {
	Msg      *dns.Msg
	Upstream string
	RTT      time.Duration
	Attempts int
	TCP      bool
}

// Resolve forwards a query upstream and returns the first usable answer.
//
// The request is copied before being modified, so the caller's message is never
// mutated.
func (r *Resolver) Resolve(ctx context.Context, req *dns.Msg) (*Result, error) {
	cfg := r.Config()
	ups := r.order()
	if len(ups) == 0 {
		return nil, ErrNoUpstreams
	}

	if err := r.acquire(ctx); err != nil {
		return nil, err
	}
	defer r.release()

	out := req.Copy()
	out.Id = dns.Id()
	out.RecursionDesired = true
	r.applyEDNS(out, cfg)

	attempts := cfg.Retries + 1
	if attempts > len(ups) {
		attempts = len(ups)
	}
	if attempts < 1 {
		attempts = 1
	}

	var lastErr error
	for i := 0; i < attempts; i++ {
		u := ups[i%len(ups)]
		res, err := r.exchange(ctx, u, out, cfg)
		if err != nil {
			lastErr = err
			u.recordFailure(err)
			r.failures.Add(1)
			if ctx.Err() != nil {
				break
			}
			continue
		}
		u.recordSuccess(res.RTT)
		r.queries.Add(1)
		prev := r.latencyUS.Load()
		r.latencyUS.Store((prev*3 + res.RTT.Microseconds()) / 4)
		res.Attempts = i + 1
		return res, nil
	}

	if lastErr == nil {
		lastErr = ErrAllFailed
	}
	return nil, fmt.Errorf("%w: %v", ErrAllFailed, lastErr)
}

// exchange performs one upstream query, falling back to TCP when the UDP
// answer comes back truncated.
func (r *Resolver) exchange(ctx context.Context, u *Upstream, req *dns.Msg, cfg Config) (*Result, error) {
	qctx, cancel := context.WithTimeout(ctx, cfg.Timeout)
	defer cancel()

	msg, rtt, err := r.udpClient.ExchangeContext(qctx, req, u.Addr)
	if err != nil {
		return nil, fmt.Errorf("query %s over UDP: %w", u.Addr, err)
	}
	if msg.Truncated {
		r.truncated.Add(1)
		tctx, tcancel := context.WithTimeout(ctx, cfg.Timeout)
		defer tcancel()
		tmsg, trtt, terr := r.tcpClient.ExchangeContext(tctx, req, u.Addr)
		if terr != nil {
			// The truncated UDP answer is still better than nothing.
			return &Result{Msg: msg, Upstream: u.Addr, RTT: rtt}, nil
		}
		return &Result{Msg: tmsg, Upstream: u.Addr, RTT: rtt + trtt, TCP: true}, nil
	}
	return &Result{Msg: msg, Upstream: u.Addr, RTT: rtt}, nil
}

// applyEDNS attaches our own OPT record, replacing whatever the client sent.
// The client's advertised buffer size describes its link, not ours.
func (r *Resolver) applyEDNS(m *dns.Msg, cfg Config) {
	m.Extra = stripOPT(m.Extra)
	opt := &dns.OPT{Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT}}
	opt.SetUDPSize(cfg.EDNSUDPSize)
	if cfg.DNSSEC {
		opt.SetDo(true)
	}
	m.Extra = append(m.Extra, opt)
}

func stripOPT(rrs []dns.RR) []dns.RR {
	out := make([]dns.RR, 0, len(rrs))
	for _, rr := range rrs {
		if rr.Header().Rrtype == dns.TypeOPT {
			continue
		}
		out = append(out, rr)
	}
	return out
}

func (r *Resolver) acquire(ctx context.Context) error {
	r.semMu.Lock()
	sem := r.sem
	r.semMu.Unlock()
	select {
	case sem <- struct{}{}:
		return nil
	case <-ctx.Done():
		return fmt.Errorf("resolver is at its concurrency limit: %w", ctx.Err())
	}
}

func (r *Resolver) release() {
	r.semMu.Lock()
	sem := r.sem
	r.semMu.Unlock()
	select {
	case <-sem:
	default:
	}
}

// Stats summarises resolver activity.
type Stats struct {
	Queries      int64   `json:"queries"`
	Failures     int64   `json:"failures"`
	Truncated    int64   `json:"truncated"`
	AvgLatencyMS float64 `json:"avg_latency_ms"`
	Upstreams    int     `json:"upstreams"`
	Healthy      int     `json:"healthy"`
}

// Stats returns the resolver counters.
func (r *Resolver) Stats() Stats {
	s := Stats{
		Queries:      r.queries.Load(),
		Failures:     r.failures.Load(),
		Truncated:    r.truncated.Load(),
		AvgLatencyMS: float64(r.latencyUS.Load()) / 1000,
	}
	for _, u := range r.Upstreams() {
		s.Upstreams++
		if u.healthy() {
			s.Healthy++
		}
	}
	return s
}

// Check performs a one-off probe against a single upstream, used by the
// "test resolver" button in the UI.
func Check(ctx context.Context, addr, qname string, timeout time.Duration) (time.Duration, string, error) {
	c := &dns.Client{Net: "udp", Timeout: timeout}
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(qname), dns.TypeA)
	m.RecursionDesired = true
	m.SetEdns0(1232, true)

	qctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	resp, rtt, err := c.ExchangeContext(qctx, m, addr)
	if err != nil {
		return 0, "", fmt.Errorf("%s did not answer: %w", addr, err)
	}
	return rtt, dns.RcodeToString[resp.Rcode], nil
}