package message

import (
	"sync"
	"time"
)

// rateLimiter 是每端一个令牌桶：速率 rps、容量 burst。
// rps<=0 表示不限速。
type rateLimiter struct {
	rps   float64
	burst float64

	mu sync.Mutex
	m  map[string]*tokenBucket
}

type tokenBucket struct {
	tokens float64
	last   time.Time
}

func newRateLimiter(rps float64, burst int) *rateLimiter {
	if burst <= 0 {
		burst = defaultRequestBurst
	}
	return &rateLimiter{
		rps:   rps,
		burst: float64(burst),
		m:     make(map[string]*tokenBucket),
	}
}

// allow 消耗 1 个令牌；允许则 true。
func (r *rateLimiter) allow(endpointID string, now time.Time) bool {
	if r == nil || r.rps <= 0 {
		return true
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	b := r.m[endpointID]
	if b == nil {
		b = &tokenBucket{tokens: r.burst, last: now}
		r.m[endpointID] = b
	}
	elapsed := now.Sub(b.last).Seconds()
	if elapsed > 0 {
		b.tokens += elapsed * r.rps
		if b.tokens > r.burst {
			b.tokens = r.burst
		}
		b.last = now
	}
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

// AllowRequest 消耗该端 1 个请求令牌；允许则 true。rps<=0 时不限速。
func (a *App) AllowRequest(endpointID string) bool {
	if a == nil {
		return true
	}
	return a.rates.allow(endpointID, a.now())
}