-- The script encapsulates the logic of a token bucket rate limiter.
-- Two types of operations are supported: 'check-only' and 'use-if-available' (controlled by the 'useTokens' arg).
-- Both operations take in rate limiter configuration parameters and the requested amount of tokens.
-- Both operations return 0, if the rate limiter has enough tokens to cover the requested amount,
-- and the deficit amount otherwise.
-- However, 'check-only' operation doesn't modify the bucket, while 'use-if-available' (if successful)
-- reduces the amount of available tokens by the requested amount.

local bucketId = KEYS[1]

local bucketSize = tonumber(ARGV[1])
local refillRatePerMillis = tonumber(ARGV[2])
local currentTimeMillis = tonumber(ARGV[3])
local requestedAmount = tonumber(ARGV[4])
local useTokens = ARGV[5] and string.lower(ARGV[5]) == "true"

local SIZE_FIELD = "s"
local TIME_FIELD = "t"

local changesMade = false
local tokensRemaining
local lastUpdateTimeMillis

-- while we're migrating from json to redis list key types, there are three possible options for the
-- type of the `bucketId` key: "string" (legacy, json value), "list" (new format), "none" (key not set).
--
-- On a separate note -- the reason we're not using a different key is because Redis Lua requires to list all keys
-- as a script input and we don't want to expose this migration to the script users.
--
-- Finally, it's okay to read the "ok" key of the return here because "TYPE" command always succeeds.
local keyType = redis.call("TYPE", bucketId)["ok"]
if keyType == "none" then
    -- if the key is not set, building the object from the configuration
    tokensRemaining = bucketSize
    lastUpdateTimeMillis = currentTimeMillis
elseif keyType == "string" then
    -- if the key is "string", we parse the value from json
    local fromJson = cjson.decode(redis.call("GET", bucketId))
    tokensRemaining = fromJson.spaceRemaining
    lastUpdateTimeMillis = fromJson.lastUpdateTimeMillis
    redis.call("DEL", bucketId)
    changesMade = true
elseif keyType == "hash" then
    -- finally, reading values from the new storage format
    local tokensRemainingStr, lastUpdateTimeMillisStr = unpack(redis.call("HMGET", bucketId, SIZE_FIELD, TIME_FIELD))
    tokensRemaining = tonumber(tokensRemainingStr)
    lastUpdateTimeMillis = tonumber(lastUpdateTimeMillisStr)
end

local elapsedTime = currentTimeMillis - lastUpdateTimeMillis
local availableAmount = math.min(
    bucketSize,
    math.floor(tokensRemaining + (elapsedTime * refillRatePerMillis))
)

if availableAmount >= requestedAmount then
    if useTokens then
        tokensRemaining = availableAmount - requestedAmount
        lastUpdateTimeMillis = currentTimeMillis
        changesMade = true
    end
    if changesMade then
        local tokensUsed = bucketSize - tokensRemaining
        -- Storing a 'full' bucket (i.e. tokensUsed == 0) is equivalent of not storing any state at all
        -- (in which case a bucket will be just initialized from the input configs as a 'full' one).
        -- For this reason, we either set an expiration time on the record (calculated to let the bucket fully replenish)
        -- or we just delete the key if the bucket is full.
        if tokensUsed > 0 then
            local ttlMillis = math.ceil(tokensUsed / refillRatePerMillis)
            redis.call("HSET", bucketId, SIZE_FIELD, tokensRemaining, TIME_FIELD, lastUpdateTimeMillis)
            redis.call("PEXPIRE", bucketId, ttlMillis)
        else
            redis.call("DEL", bucketId)
        end
    end
    return 0
else
    return requestedAmount - availableAmount
end