// Package resil implements the data-plane resilience toolkit: power-of-
// two-choices load balancing with EWMA latency and outlier ejection,
// adaptive concurrency limiting (Little's Law / gradient), retry budgets,
// and a circuit breaker. Stdlib-only and deterministic (seeded RNG,
// injectable clock) so the behaviors are testable.
package resil

import (
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

// Endpoint is one upstream instance with live load/health signals.
type Endpoint struct {
	Addr string

	inflight   atomic.Int64
	ewmaNs     atomic.Int64 // EWMA of observed latency, nanoseconds
	fails      atomic.Int64 // consecutive failures
	ejectUntil atomic.Int64 // unix-nano; 0 = healthy
}

// score: lower is better. Combines in-flight load and latency EWMA so
// P2C routes toward the less-loaded, faster endpoint.
func (e *Endpoint) score() float64 {
	return float64(e.inflight.Load()+1) * float64(e.ewmaNs.Load()+1)
}

func (e *Endpoint) ejected(nowNs int64) bool { return e.ejectUntil.Load() > nowNs }

// P2C is a power-of-two-choices balancer with outlier ejection.
type P2C struct {
	mu  sync.Mutex
	eps []*Endpoint
	rng *rand.Rand
	now func() time.Time

	EjectThreshold int           // consecutive failures before ejection
	BaseCooldown   time.Duration // ejection cool-down (grows with repeats)
	PanicFraction  float64       // if more than this fraction would be ejected, eject none
}

func NewP2C(addrs []string, seed int64) *P2C {
	b := &P2C{
		rng:            rand.New(rand.NewSource(seed)),
		now:            time.Now,
		EjectThreshold: 5,
		BaseCooldown:   2 * time.Second,
		PanicFraction:  0.5,
	}
	for _, a := range addrs {
		b.eps = append(b.eps, &Endpoint{Addr: a})
	}
	return b
}

func (b *P2C) Endpoints() []*Endpoint { return b.eps }

// Pick returns an endpoint via power-of-two-choices over the healthy set.
// If more than PanicFraction of endpoints are ejected, it ignores
// ejection (panic mode) so a fleet-wide brown-out never black-holes all
// traffic.
func (b *P2C) Pick() *Endpoint {
	b.mu.Lock()
	defer b.mu.Unlock()
	nowNs := b.now().UnixNano()

	healthy := make([]*Endpoint, 0, len(b.eps))
	for _, e := range b.eps {
		if !e.ejected(nowNs) {
			healthy = append(healthy, e)
		}
	}
	ejected := len(b.eps) - len(healthy)
	if len(healthy) == 0 || float64(ejected) > b.PanicFraction*float64(len(b.eps)) {
		healthy = b.eps // panic mode
	}
	if len(healthy) == 1 {
		return healthy[0]
	}
	i := b.rng.Intn(len(healthy))
	j := b.rng.Intn(len(healthy) - 1)
	if j >= i {
		j++ // distinct second choice
	}
	a, c := healthy[i], healthy[j]
	if a.score() <= c.score() {
		return a
	}
	return c
}

// Acquire/Release track in-flight load for the score (call around a request).
func (b *P2C) Acquire(e *Endpoint) { e.inflight.Add(1) }
func (b *P2C) Release(e *Endpoint) { e.inflight.Add(-1) }

// Observe records a request result: updates the latency EWMA and drives
// consecutive-failure ejection with a growing cool-down.
func (b *P2C) Observe(e *Endpoint, latency time.Duration, ok bool) {
	const alpha = 0.2
	old := e.ewmaNs.Load()
	sample := latency.Nanoseconds()
	e.ewmaNs.Store(int64(alpha*float64(sample) + (1-alpha)*float64(old)))

	if ok {
		e.fails.Store(0)
		return
	}
	n := e.fails.Add(1)
	if int(n) >= b.EjectThreshold {
		cool := time.Duration(n) * b.BaseCooldown
		e.ejectUntil.Store(b.now().Add(cool).UnixNano())
	}
}