--[[
    配额令牌桶算法实现
    
    功能说明:
    - 实现基于配额的限流
    - 支持TTL自动重置配额
    - 适用于固定时间窗口内的配额管理
    
    参数说明:
    - KEYS[1]: 令牌桶key
    - ARGV[1]: capacity - 令牌桶容量(最大配额)
    - ARGV[2]: required - 需要获取的令牌数
    - ARGV[3]: ttl - 令牌桶过期时间(秒)
    
    返回值:
    - {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 capacity = tonumber(ARGV[1])
local required = tonumber(ARGV[2])
local ttl = tonumber(ARGV[3])
local ttlms = ttl * 1000

-- Hash字段常量
local TOKENS_FIELD = "tokens"
local EXPIRE_TIME_FIELD = "expire_time"
local CAPACITY_FIELD = "capacity"

local current_ms = getCurrentMillis()
local ret = {}

-- 检查key是否存在
local exists = redis.call('EXISTS', key)

if exists == 1 then
    -- 使用HMGET直接获取指定字段
    local state = redis.call('HMGET', key, TOKENS_FIELD, EXPIRE_TIME_FIELD)
    local current_tokens = tonumber(state[1]) or capacity
    local expire_time = tonumber(state[2]) or (current_ms + ttlms)
    
    -- 检查是否过期,过期则重置配额
    if current_ms > expire_time then
        current_tokens = capacity
        expire_time = current_ms + ttlms
    end
    
    -- 检查令牌是否足够
    if required > current_tokens then
        return {0, current_tokens}
    end
    
    -- 消耗令牌
    local acquired_tokens = math.min(required, current_tokens)
    current_tokens = current_tokens - acquired_tokens
    
    -- 更新状态
    redis.call('HMSET', key, 
        TOKENS_FIELD, current_tokens, 
        CAPACITY_FIELD, capacity, 
        EXPIRE_TIME_FIELD, expire_time)
    redis.call('EXPIRE', key, math.max(ttl, 60))
    
    ret = {acquired_tokens, current_tokens}
else
    -- 初始化令牌桶
    if required > capacity then
        return {0, capacity}
    end
    
    local remaining_tokens = math.max(0, capacity - required)
    local new_expire_time = current_ms + ttlms
    
    redis.call('HMSET', key,
        TOKENS_FIELD, remaining_tokens,
        CAPACITY_FIELD, capacity,
        EXPIRE_TIME_FIELD, new_expire_time)
    redis.call('EXPIRE', key, math.max(ttl, 60))
    
    ret = {required, remaining_tokens}
end

return ret