package limit

import (
	"context"
	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
	"net/http"
	"sync"
	"time"
)

var ipRateLimiter *IPRateLimiter
var config *Config

type Config struct {
	Rate     int
	Bursts   int
	WaitTime int
}

type (
	Option func(c *Config)
)

func Rate(rate int) Option {
	return func(c *Config) {
		c.Rate = rate
	}
}

func Bursts(bursts int) Option {
	return func(c *Config) {
		c.Bursts = bursts
	}
}

func WaitTime(waitTime int) Option {
	return func(c *Config) {
		c.WaitTime = waitTime
	}
}

func SetUp(options ...Option) {
	config = &Config{
		Rate:     1,
		Bursts:   60,
		WaitTime: 10,
	}
	for _, option := range options {
		option(config)
	}
}

// IPRateLimiter .
type IPRateLimiter struct {
	ips map[string]*rate.Limiter
	mu  *sync.RWMutex
}

// NewIPRateLimiter .
func NewIPRateLimiter() *IPRateLimiter {
	ipRateLimiter = &IPRateLimiter{
		ips: make(map[string]*rate.Limiter),
		mu:  &sync.RWMutex{},
	}
	return ipRateLimiter
}

func GetIPRateLimiter() *IPRateLimiter {
	if ipRateLimiter == nil {
		NewIPRateLimiter()
	}
	return ipRateLimiter
}

// AddIP 创建了一个新的速率限制器，并将其添加到 ips 映射中,
// 使用 IP地址作为密钥
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
	i.mu.Lock()
	defer i.mu.Unlock()

	limit := rate.Every(time.Duration(config.Rate) * time.Second)
	limiter := rate.NewLimiter(limit, config.Bursts)
	i.ips[ip] = limiter

	return limiter
}

// GetLimiter 返回所提供的IP地址的速率限制器(如果存在的话).
// 否则调用 AddIP 将 IP 地址添加到映射中
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
	i.mu.Lock()
	limiter, exists := i.ips[ip]

	if !exists {
		i.mu.Unlock()
		return i.AddIP(ip)
	}

	i.mu.Unlock()

	return limiter
}

func Limiter() gin.HandlerFunc {
	return func(c *gin.Context) {
		limiter := GetIPRateLimiter().GetLimiter(c.ClientIP())
		if limiter.Allow() {
			c.Next()
		} else {
			ctx, cancel := context.WithTimeout(context.TODO(), time.Duration(config.WaitTime)*time.Second)
			defer cancel()
			err := limiter.Wait(ctx)
			if err == nil {
				c.Next()
			} else {
				c.JSON(http.StatusTooManyRequests, gin.H{
					"code": http.StatusTooManyRequests,
					"msg":  "to many requests",
					"data": "to many requests",
				})
				c.Abort()
			}
		}
	}
}