package ratelimit

import (
	"sync"
	"time"
)

// Limiter is a per-key token bucket. Keys are RealIP (C6), never gnet RemoteAddr.
type Limiter struct {
	mu      sync.Mutex
	rate    float64
	burst   float64
	buckets map[string]*bucket
}

type bucket struct {
	tokens float64
	last   time.Time
}

func New(perSec int) *Limiter {
	if perSec < 1 {
		perSec = 5
	}
	return &Limiter{
		rate:    float64(perSec),
		burst:   float64(perSec),
		buckets: make(map[string]*bucket),
	}
}

func (l *Limiter) Allow(key string) bool {
	if l == nil {
		return true
	}
	if key == "" {
		key = "unknown"
	}
	now := time.Now()
	l.mu.Lock()
	defer l.mu.Unlock()
	b := l.buckets[key]
	if b == nil {
		b = &bucket{tokens: l.burst, last: now}
		l.buckets[key] = b
	}
	elapsed := now.Sub(b.last).Seconds()
	b.tokens += elapsed * l.rate
	if b.tokens > l.burst {
		b.tokens = l.burst
	}
	b.last = now
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	if len(l.buckets) > 8192 {
		l.gc(now)
	}
	return true
}

func (l *Limiter) gc(now time.Time) {
	for k, b := range l.buckets {
		if now.Sub(b.last) > time.Minute {
			delete(l.buckets, k)
		}
	}
}