--[[
    QPS令牌桶算法实现
    
    功能说明:
    - 实现基于令牌桶的QPS限流
    - 支持突发流量处理(burst)
    - 自动补充令牌
    
    参数说明:
    - KEYS[1]: 令牌桶key
    - ARGV[1]: token_interval_ms - 生产一个令牌需要的毫秒数
    - ARGV[2]: burst_limit - 令牌桶最大容量(突发上限)
    - ARGV[3]: required_tokens - 需要获取的令牌数目
    
    返回值:
    - {acquired_tokens, remaining_tokens}
    - acquired_tokens: 实际获取的令牌数(0表示获取失败)
    - remaining_tokens: 获取后剩余的令牌数
]]

redis.replicate_commands()

-- 获取当前毫秒时间戳
local function getCurrentMillis()
    local current = redis.call('TIME')
    return current[1] * 1000 + math.modf(current[2] / 1000)
end

-- 获取输入参数
local key = KEYS[1]
local token_interval_ms = tonumber(ARGV[1])
local burst_limit = tonumber(ARGV[2])
local required_tokens = tonumber(ARGV[3])
local current_time_ms = getCurrentMillis()

-- Hash字段常量
local TOKENS_FIELD = "tokens"
local LAST_REFILL_FIELD = "last_refill"

-- 使用HMGET直接获取指定字段,避免HGETALL的解析开销
local state = redis.call('HMGET', key, TOKENS_FIELD, LAST_REFILL_FIELD)
local current_tokens, last_refill

if not state[1] then
    current_tokens = burst_limit
    last_refill = current_time_ms
else
    current_tokens = tonumber(state[1]) or burst_limit
    last_refill = tonumber(state[2]) or current_time_ms
end

-- 计算自上次补充以来新增加的令牌数
local elapsed_ms = current_time_ms - last_refill
local new_tokens = elapsed_ms / token_interval_ms
local updated_tokens = math.min(burst_limit, current_tokens + new_tokens)

-- 计算实际可获取的令牌数

local acquired_tokens = math.min(required_tokens, updated_tokens)

-- 令牌不足,直接返回,不更新状态
if acquired_tokens < 1 then
    return {0, updated_tokens}
end

-- 消耗令牌
updated_tokens = updated_tokens - acquired_tokens

-- 计算过期时间: 至少存储一个完整令牌桶填充周期
local expire = math.ceil(burst_limit * token_interval_ms / 1000) + 1
expire = math.max(expire, 60)

-- 更新令牌桶状态
redis.call('HMSET', key, TOKENS_FIELD, updated_tokens, LAST_REFILL_FIELD, current_time_ms)
redis.call('EXPIRE', key, expire)

return {acquired_tokens, updated_tokens}