package ratelimit

import (
	"sync"
	"time"
)

// Limiter is a token-bucket rate limiter keyed by string.
type Limiter struct {
	mu    sync.Mutex
	rate  float64 // tokens per second
	burst float64
	tokens map[string]*bucket
}

type bucket struct {
	tokens   float64
	lastFill time.Time
}

// New creates a limiter refilling `rate` tokens/sec with `burst` capacity.
func New(rate float64, burst int) *Limiter {
	return &Limiter{
		rate:   rate,
		burst:  float64(burst),
		tokens: map[string]*bucket{},
	}
}

// Allow checks whether `key` may take one token now.
func (l *Limiter) Allow(key string) bool {
	return l.Take(key, 1)
}

// Take checks whether `key` may take n tokens now.
func (l *Limiter) Take(key string, n float64) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := time.Now()
	b, ok := l.tokens[key]
	if !ok {
		b = &bucket{tokens: l.burst, lastFill: now}
		l.tokens[key] = b
	}
	// Refill based on elapsed time.
	elapsed := now.Sub(b.lastFill).Seconds()
	b.tokens = minF(l.burst, b.tokens+elapsed*l.rate)
	b.lastFill = now
	if b.tokens >= n {
		b.tokens -= n
		return true
	}
	return false
}

// Sweep removes idle buckets to bound memory. Call periodically.
func (l *Limiter) Sweep(olderThan time.Duration) {
	l.mu.Lock()
	defer l.mu.Unlock()
	cutoff := time.Now().Add(-olderThan)
	for k, b := range l.tokens {
		if b.lastFill.Before(cutoff) {
			delete(l.tokens, k)
		}
	}
}

func minF(a, b float64) float64 {
	if a < b {
		return a
	}
	return b
}