// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

// Package ratelimit implements the spec §2.5 token buckets: per-credential and per-source-IP
// ceilings on session creation, actions, UDP packets and bytes. One Limiter holds one policy
// (rate + burst) and lazily creates a bucket per key; callers namespace their keys ("cred:…",
// "ip:…") so a single Limiter can enforce both axes of the same rule.
package ratelimit

import (
	"sync"
	"time"
)

// Limiter is a keyed set of token buckets sharing one rate and burst.
//
// A nil *Limiter allows everything: the ceilings are configurable down to "off" (config value 0),
// and a nil check in one place beats a sentinel policy that every call site must know about.
type Limiter struct {
	rate  float64 // tokens per second
	burst float64

	mu        sync.Mutex
	buckets   map[string]*bucket
	lastSweep time.Time
	now       func() time.Time // swappable so tests need no sleeping
}

type bucket struct {
	tokens float64
	last   time.Time
}

// New creates a limiter granting ratePerSec tokens per second per key, holding at most burst.
func New(ratePerSec, burst float64) *Limiter {
	return &Limiter{
		rate:    ratePerSec,
		burst:   burst,
		buckets: map[string]*bucket{},
		now:     time.Now,
	}
}

// Allow takes one token for key. See AllowN.
func (l *Limiter) Allow(key string) (bool, time.Duration) {
	return l.AllowN(key, 1)
}

// AllowN takes n tokens for key, reporting whether they were available and — when they were
// not — how long until they will be, which is what a control-plane 429 puts in Retry-After.
// A refusal consumes nothing: the caller being told to wait must not itself push the wait out.
func (l *Limiter) AllowN(key string, n float64) (bool, time.Duration) {
	if l == nil {
		return true, 0
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	now := l.now()
	l.sweepLocked(now)
	b := l.buckets[key]
	if b == nil {
		b = &bucket{tokens: l.burst, last: now}
		l.buckets[key] = b
	}
	b.tokens += now.Sub(b.last).Seconds() * l.rate
	if b.tokens > l.burst {
		b.tokens = l.burst
	}
	b.last = now
	if b.tokens >= n {
		b.tokens -= n
		return true, 0
	}
	return false, time.Duration((n - b.tokens) / l.rate * float64(time.Second))
}

// sweepEvery bounds how often the map is walked; the walk is cheap but there is no point doing
// it per packet on the data plane's hot path.
const sweepEvery = time.Minute

// sweepLocked drops buckets that have been idle long enough to be full again. A full bucket
// carries no state a fresh one would not, and without the sweep the map grows one entry per
// source address ever seen — an attacker-controlled key space must not be an unbounded one.
func (l *Limiter) sweepLocked(now time.Time) {
	if now.Sub(l.lastSweep) < sweepEvery {
		return
	}
	l.lastSweep = now
	idle := sweepEvery
	if l.rate > 0 {
		if refill := time.Duration(l.burst / l.rate * float64(time.Second)); refill > idle {
			idle = refill
		}
	}
	for k, b := range l.buckets {
		if now.Sub(b.last) > idle {
			delete(l.buckets, k)
		}
	}
}