package middleware

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

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

type visitorEntry struct {
	limiter *rate.Limiter
	last    time.Time
}

// RateLimit 按 ClientIP 的令牌桶限流中间件 // Logan.Li
func RateLimit(rps float64, burst int) gin.HandlerFunc {
	var mu sync.Mutex
	visitors := map[string]*visitorEntry{}

	go func() {
		ticker := time.NewTicker(5 * time.Minute)
		defer ticker.Stop()
		for range ticker.C {
			mu.Lock()
			cutoff := time.Now().Add(-10 * time.Minute)
			for ip, e := range visitors {
				if e.last.Before(cutoff) {
					delete(visitors, ip)
				}
			}
			mu.Unlock()
		}
	}()

	return func(c *gin.Context) {
		ip := c.ClientIP()
		mu.Lock()
		e, ok := visitors[ip]
		if !ok {
			e = &visitorEntry{limiter: rate.NewLimiter(rate.Limit(rps), burst)}
			visitors[ip] = e
		}
		e.last = time.Now()
		allow := e.limiter.Allow()
		mu.Unlock()
		if !allow {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"code":    429,
				"message": "请求过于频繁，请稍后再试",
			})
			return
		}
		c.Next()
	}
}