package limit

import (
	"fmt"
	"stat/common/redis"
	"stat/statg"
	"time"

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

func GetKey(name string) func() string {
	return func() string {
		return redis.GetIpLimitDayKey(statg.Conf.Host, statg.Conf.Port) + "_" + name
	}
}

func Limit(name string, capacity int64, rate int64) gin.HandlerFunc {
	limiter := redis.NewTokenBucket(redis.StatRds.Client(), GetKey(name), capacity, rate, 48*time.Hour)
	return func(ctx *gin.Context) {
		if err := limiter.TakeTokens(ctx.ClientIP(), 1); err != nil {
			ctx.AbortWithStatusJSON(200, gin.H{"msg": fmt.Sprintf("%+v", err)})
		}
	}
}