package middleware

import (
	"server/common"
	"server/db"
	"server/modules/web/app"
	"time"

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

const (
	UserLimitPerSec = 10
)

func RateLimitMiddleware() gin.HandlerFunc {
	l := rate.NewLimiter(5000, 20000)
	return func(c *gin.Context) {
		if !l.Allow() {
			c.Abort()
			return
		}
		c.Next()
	}
}

func UserLimitMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		a := app.NewApp(c)
		ip := a.GetRemoteIP()
		key := common.GetRedisLockKeyIP(ip)
		if db.Redis().Exist(key) {
			res, err := db.Redis().Incr(key, 1)
			if err != nil {
				log.Error("err:%v", err)
				c.Abort()
				return
			}
			if res > UserLimitPerSec {
				c.Abort()
				return
			}
		} else {
			db.Redis().SetData(key, 0, time.Second)
		}
		c.Next()
	}
}