package middleware

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/store/memory"
)

// RateLimitConfig 限流配置
type RateLimitConfig struct {
	GlobalRPS  float64
	PerIPRPS   float64
	PerIPBurst int
}

// RateLimiter 双维度限流：全局令牌桶 + 单 IP 令牌桶
// 开发/测试环境可用内存存储，生产环境建议换 Redis 存储
func RateLimiter(cfg RateLimitConfig) gin.HandlerFunc {
	store := memory.NewStore()

	// 全局限流
	globalRate := limiter.Rate{
		Period: time.Second,
		Limit:  int64(cfg.GlobalRPS),
	}
	globalLimiter := limiter.New(store, globalRate)

	// 单 IP 限流
	perIPRate := limiter.Rate{
		Period: time.Second,
		Limit:  int64(cfg.PerIPRPS),
	}
	perIPLimiter := limiter.New(store, perIPRate)

	return func(c *gin.Context) {
		// 1. 检查全局限流
		globalCtx, err := globalLimiter.Get(c.Request.Context(), "global")
		if err != nil || globalCtx.Reached {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    4290,
				"message": "server busy, please retry later",
			})
			c.Abort()
			return
		}

		// 2. 检查单 IP 限流
		ip := c.ClientIP()
		ipCtx, err := perIPLimiter.Get(c.Request.Context(), "ip:"+ip)
		if err != nil || ipCtx.Reached {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    4291,
				"message": "too many requests from this IP",
			})
			c.Abort()
			return
		}

		c.Next()
	}
}