// Package middleware ...
//
// Description : 基于redis的流控管理
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021-07-30 20:52 下午
package middleware

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

	"github.com/pkg/errors"

	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"github.com/go-redis/redis_rate/v9"
)

// BuildRateLimitConfig 构建流控配置
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:35 下午 2021/7/30
func BuildRateLimitConfig(rate int, burst int, period time.Duration) redis_rate.Limit {
	return redis_rate.Limit{
		Rate:   rate,
		Burst:  burst,
		Period: period,
	}
}

// RateLimit gin 可以直接引用的中间件
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 10:35 下午 2021/7/30
func RateLimit(redisInstance *redis.Client, key string, limitRule redis_rate.Limit, responseData map[string]interface{}) gin.HandlerFunc {
	return func(ctx *gin.Context) {
		var (
			err error
		)

		if _, err = Allow(redisInstance, key, limitRule); nil != err {
			if nil == responseData {
				responseData = make(map[string]interface{})
				responseData["code"] = -1
				responseData["message"] = err.Error()
				responseData["data"] = make(map[string]interface{})
			}
			ctx.AbortWithStatusJSON(http.StatusOK, responseData)
			return
		}

		ctx.Next()
	}
}

// Allow 是否允许访问
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 11:52 下午 2021/7/30
func Allow(redisInstance *redis.Client, key string, limitRule redis_rate.Limit) (bool, error) {
	var (
		limiter *redis_rate.Limiter
		result  *redis_rate.Result
		err     error
	)

	limiter = redis_rate.NewLimiter(redisInstance)
	if result, err = limiter.Allow(context.Background(), key, limitRule); nil != err {
		return false, err
	}

	if result.Allowed > 0 {
		return true, nil
	}

	return false, errors.New("rate limit! current request is not allowed!")
}