type FixedWindowRateLimiter struct {
    threshold int           // 阈值
    stime     time.Time     // 开始时间
    interval  time.Duration // 时间窗口
    counter   int           // 当前计数
    lock      sync.Mutex
}

func NewFixedWindowRateLimiter(threshold int, interval time.Duration) *FixedWindowRateLimiter {

    return &FixedWindowRateLimiter{
        threshold: threshold,
        stime:     time.Now(),
        interval:  interval,
        counter:   threshold - 1, // 让其处于下一个时间窗口开始的时间临界点
    }
}

func (l *FixedWindowRateLimiter) Allow() bool {
    l.lock.Lock()
    defer l.lock.Unlock()

    // 判断收到请求数是否达到阈值
    if l.counter == l.threshold-1 {
        now := time.Now()
        // 达到阈值后，判断是否是请求窗口内
        if now.Sub(l.stime) >= l.interval {
            // 重新计数
            l.Reset()
            return true
        }
        // 丢弃多余的请求
        return false
    } else {
        l.counter++
        return true
    }
}

func (l *FixedWindowRateLimiter) Reset() {
    l.counter = 0
    l.stime = time.Now()
}

type LeakyBucketRateLimiter struct {
    capacity   float64 `json:"capacity"`   // 桶的容量
    water      float64 `json:"water"`      // 当前桶中水量
    flowRate   float64 `json:"flowRate"`   // 每秒漏桶流速
    lastLeakMs int64   `json:"lastLeakMs"` // 上次漏水毫秒数
    lock       sync.Mutex
}

func NewLeakyBucketRateLimiter(flowRate, capacity float64) *LeakyBucketRateLimiter {
    return &LeakyBucketRateLimiter{
        capacity:   capacity,
        flowRate:   flowRate,
        water:      capacity + 1, // 设置起始边界
        lastLeakMs: time.Now().UnixNano() / 1e6,
    }
}

func (l *LeakyBucketRateLimiter) Allow() bool {
    l.lock.Lock()
    defer l.lock.Unlock()

    // 获取当前时间
    now := time.Now().UnixNano() / 1e6
    // 计算这段时间流出的水量：
    outflowWater := (float64(now - l.lastLeakMs)) * l.flowRate / 1000
    // 计算水量： 桶的当前水量 - 流出的水量
    l.water = math.Max(0, l.water-outflowWater)
    l.lastLeakMs = now
    if l.water < l.capacity {
        // 当前水量 小于 桶容量，允许通过
        l.water++
        return true
    } else {
        // 当前水量 不小于 桶容量，不允许通过
        return false
    }
}

-- 计数器限流
-- 此处支持的最小单位时间是秒, 若将 expire 改成 pexpire 则可支持毫秒粒度.
-- KEYS[1]  string  限流的key
-- ARGV[1]  int     限流数
-- ARGV[2]  int     单位时间(秒)

local cnt = tonumber(redis.call("incr", KEYS[1]))

if (cnt == 1) then
    -- cnt 值为1说明之前不存在该值, 因此需要设置其过期时间
    redis.call("expire", KEYS[1], tonumber(ARGV[2]))
elseif (cnt > tonumber(ARGV[1])) then
    return -1
end 

return cnt

-- 令牌桶限流: 不支持预消费, 初始桶是满的
-- KEYS[1]  string  限流的key

-- ARGV[1]  int     桶最大容量
-- ARGV[2]  int     每次添加令牌数
-- ARGV[3]  int     令牌添加间隔(秒)
-- ARGV[4]  int     当前时间戳

local bucket_capacity = tonumber(ARGV[1])
local add_token = tonumber(ARGV[2])
local add_interval = tonumber(ARGV[3])
local now = tonumber(ARGV[4])

-- 保存上一次更新桶的时间的key
local LAST_TIME_KEY = KEYS[1].."_time";         
-- 获取当前桶中令牌数
local token_cnt = redis.call("get", KEYS[1])    
-- 桶完全恢复需要的最大时长
local reset_time = math.ceil(bucket_capacity / add_token) * add_interval;

if token_cnt then   -- 令牌桶存在
    -- 上一次更新桶的时间
    local last_time = redis.call('get', LAST_TIME_KEY)
    -- 恢复倍数
    local multiple = math.floor((now - last_time) / add_interval)
    -- 恢复令牌数
    local recovery_cnt = multiple * add_token
    -- 确保不超过桶容量
    local token_cnt = math.min(bucket_capacity, token_cnt + recovery_cnt) - 1
    
    if token_cnt < 0 then
        return -1;
    end
    
    -- 重新设置过期时间, 避免key过期
    redis.call('set', KEYS[1], token_cnt, 'EX', reset_time)                     
    redis.call('set', LAST_TIME_KEY, last_time + multiple * add_interval, 'EX', reset_time)
    return token_cnt
    
else    -- 令牌桶不存在
    token_cnt = bucket_capacity - 1
    -- 设置过期时间避免key一直存在
    redis.call('set', KEYS[1], token_cnt, 'EX', reset_time);
    redis.call('set', LAST_TIME_KEY, now, 'EX', reset_time + 1);    
    return token_cnt    
end