package ratelimit

import (
	"context"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/ratelimit"
)

var rl = ratelimit.New(20, ratelimit.WithoutSlack)

func RateLimitMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 1. 给请求设置超时（如3秒，超过3秒未拿到令牌直接返回）
		ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second)
		defer cancel()

		// 2. 用带超时的方式获取令牌（需自定义限流器的 Take 逻辑）
		ch := make(chan struct{}, 1)
		go func() {
			rl.Take()        // 阻塞拿令牌
			ch <- struct{}{} // 拿到令牌后通知主协程
		}()

		select {
		case <-ch:
			// 3. 拿到令牌，继续执行
			c.Next()
		case <-ctx.Done():
			// 4. 超时未拿到令牌，返回请求频繁
			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    429,
				"message": "请求过于频繁，请稍后重试",
				"reason":  "等待令牌超时",
			})
			c.Abort()
		}
	}
}