package middleware

import (
	"net/http"
	"sync"
	"time"

	"golang.org/x/time/rate"

	"github.com/gin-gonic/gin"
)

// RateLimiter 按 IP 的限流器
type RateLimiter struct {
	mu      sync.Mutex
	clients map[string]*rate.Limiter
	r       rate.Limit
	b       int
}

// NewRateLimiter 创建限流中间件，r 每秒请求数，b 突发容量
func NewRateLimiter(r rate.Limit, b int) *RateLimiter {
	return &RateLimiter{
		clients: make(map[string]*rate.Limiter),
		r:       r,
		b:       b,
	}
}

// getLimiter 获取或创建该 key 的 limiter
func (rl *RateLimiter) getLimiter(key string) *rate.Limiter {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	if lim, ok := rl.clients[key]; ok {
		return lim
	}
	lim := rate.NewLimiter(rl.r, rl.b)
	rl.clients[key] = lim
	return lim
}

// Middleware 返回 Gin 限流中间件（按客户端 IP）
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		key := c.ClientIP()
		lim := rl.getLimiter(key)
		if !lim.Allow() {
			c.AbortWithStatus(http.StatusTooManyRequests)
			return
		}
		c.Next()
	}
}

// Cleanup 定期清理过期 limiter（可选，避免 map 无限增长）
func (rl *RateLimiter) Cleanup(interval time.Duration) {
	ticker := time.NewTicker(interval)
	go func() {
		for range ticker.C {
			rl.mu.Lock()
			rl.clients = make(map[string]*rate.Limiter)
			rl.mu.Unlock()
		}
	}()
}