package limit

import (
	"fmt"
	"golang.org/x/time/rate"
	"sync"
	"time"
)

type RateLimiter struct {
	// RateLimiterMap 用来存储每个 key 对应的限流器
	rateLimiters map[string]*rate.Limiter
	mu           sync.Mutex
}

type Parameters struct {
	Key      string        // 限流器的标识
	Burst    int           // 每分钟允许的最大请求数
	Interval time.Duration // 请求数量的时间窗口内
}

func NewRateLimiter() *RateLimiter {
	return &RateLimiter{
		rateLimiters: make(map[string]*rate.Limiter),
	}
}

// RateLimit 创建一个函数来检查某个 ID 的请求是否频繁
func (r *RateLimiter) RateLimit(p *Parameters) error {
	// 锁定访问 rateLimiters，确保线程安全
	r.mu.Lock()
	defer r.mu.Unlock()

	// 如果该 id 没有对应的限流器，则创建一个新的
	if _, exists := r.rateLimiters[p.Key]; !exists {
		// 每分钟允许 1 次请求
		r.rateLimiters[p.Key] = rate.NewLimiter(rate.Every(p.Interval), p.Burst)
	}

	// 获取该 id 对应的限流器
	limiter := r.rateLimiters[p.Key]

	// 判断请求是否被限流
	if !limiter.Allow() {
		return fmt.Errorf("请求频率过快，请稍后再试")
	}

	// 请求通过限制，可以继续执行
	return nil
}
