package pushy

import "math/rand"

// Jitter generates reconnect delays. A seeded RNG keeps tests
// deterministic; production seeds from a real source.
type Jitter struct{ rng *rand.Rand }

func NewJitter(seed int64) *Jitter { return &Jitter{rng: rand.New(rand.NewSource(seed))} }

// FullJitterMillis returns a uniform random delay in [0, max) — "full
// jitter" (AWS builders' library). This is what turns a synchronized
// reconnect stampede into a smooth ramp across the surviving fleet.
func (j *Jitter) FullJitterMillis(max int) int {
	if max <= 0 {
		return 0
	}
	return j.rng.Intn(max)
}

// Backoff returns an exponential-backoff delay with full jitter, capped:
// random in [0, min(base<<attempt, cap)). Used for client reconnects
// after an UNREQUESTED disconnect.
func (j *Jitter) Backoff(attempt, baseMs, capMs int) int {
	d := baseMs << attempt
	if d > capMs || d <= 0 {
		d = capMs
	}
	return j.FullJitterMillis(d)
}

// ReconnectDelays assigns each device a reconnect delay spread uniformly
// over [0, windowMs). Sending these in a drain "reconnect" message lets
// devices migrate off this node smoothly instead of all at once.
func (j *Jitter) ReconnectDelays(devices []string, windowMs int) map[string]int {
	out := make(map[string]int, len(devices))
	for _, d := range devices {
		out[d] = j.FullJitterMillis(windowMs)
	}
	return out
}

// Dedup tracks recently-seen message IDs so a client can suppress
// duplicates from at-least-once redelivery (e.g. during a drain when a
// message is forwarded to a node the device just left). Effectively-once
// from the user's view = at-least-once delivery + idempotent client.
type Dedup struct {
	seen  map[string]struct{}
	order []string
	max   int
}

func NewDedup(max int) *Dedup {
	if max <= 0 {
		max = 1024
	}
	return &Dedup{seen: make(map[string]struct{}), max: max}
}

// Seen reports whether id was already delivered, and records it.
func (d *Dedup) Seen(id string) bool {
	if _, ok := d.seen[id]; ok {
		return true
	}
	d.seen[id] = struct{}{}
	d.order = append(d.order, id)
	if len(d.order) > d.max { // bounded memory: forget the oldest
		oldest := d.order[0]
		d.order = d.order[1:]
		delete(d.seen, oldest)
	}
	return false
}

// SpreadBuckets buckets delays into n buckets over [0, windowMs) — a test
// helper to show reconnects are spread, not stampeding.
func SpreadBuckets(delays map[string]int, windowMs, n int) []int {
	buckets := make([]int, n)
	for _, d := range delays {
		b := d * n / windowMs
		if b >= n {
			b = n - 1
		}
		buckets[b]++
	}
	return buckets
}