// Package ratelimit provides a cooperative token-bucket rate limiter.
// No mutexes — designed for single-threaded cooperative scheduling.
package ratelimit

import "time"

type bucket struct {
	tokens float64
	last   int64 // unix nano
}

// Limiter is a per-key token bucket rate limiter.
type Limiter struct {
	buckets map[string]*bucket
	rate    float64 // tokens per second
	burst   int     // max tokens
}

// New creates a rate limiter. Rate is tokens/second, burst is the
// maximum tokens that can accumulate.
func New(rate float64, burst int) *Limiter {
	return &Limiter{
		buckets: map[string]*bucket{},
		rate:    rate,
		burst:   burst,
	}
}

// Allow checks whether key has a token available. Consumes one token
// if allowed.
func (l *Limiter) Allow(key []byte) bool {
	k := string(key)
	now := time.Now().UnixNano()
	b, ok := l.buckets[k]
	if !ok {
		b = &bucket{tokens: float64(l.burst), last: now}
		l.buckets[k] = b
	}
	elapsed := float64(now-b.last) / 1e9
	b.tokens += elapsed * l.rate
	if b.tokens > float64(l.burst) {
		b.tokens = float64(l.burst)
	}
	b.last = now
	if b.tokens < 1.0 {
		return false
	}
	b.tokens--
	return true
}

// Cleanup removes entries older than maxAge to prevent unbounded growth.
func (l *Limiter) Cleanup(maxAge time.Duration) {
	cutoff := time.Now().Add(-maxAge).UnixNano()
	for k, b := range l.buckets {
		if b.last < cutoff {
			delete(l.buckets, k)
		}
	}
}