--- Lua scripts are run atomically by default, and since redis
--- is single threaded, there are no race conditions to worry about.
---
--- This script does three things, in order:
--- 1. Retrieves token bucket state, which means the last slot assigned,
---    and how many tokens are left to be assigned for that slot
--- 2. Works out whether we need to move to the next slot(s), or consume
---    tokens from the current one.
--- 3. Saves the token bucket state and returns the slot. The state is a
---    combination of the last slot assigned (timestamp) and the number of tokens left.
---
--- The token bucket implementation is forward looking, so we're really just handing
--- out the next time there would be tokens in the bucket, and letting the client
--- decide wait until that time.
---
--- returns:
--- * The assigned slot, as a millisecond timestamp

redis.replicate_commands()

-- Arguments
local capacity = tonumber(ARGV[1])
local refill_amount = tonumber(ARGV[2])
local initial_tokens = tonumber(ARGV[3])
local refill_frequency = tonumber(ARGV[4]) * 1000 -- Convert to milliseconds
local expiry = tonumber(ARGV[5])
local tokens_to_consume = tonumber(ARGV[6]) -- Number of tokens to consume
local max_sleep_ms = tonumber(ARGV[7]) * 1000 -- Convert to milliseconds
local window_start_time = tonumber(ARGV[8]) * 1000 or 0 -- Unix timestamp (0 means no alignment)

-- Validate that tokens_to_consume doesn't exceed capacity
if tokens_to_consume > capacity then
    return redis.error_reply("Requested tokens exceed bucket capacity")
end

-- Validate that tokens_to_consume is non-negative
if tokens_to_consume < 0 then
    return redis.error_reply("Can't consume negative tokens")
end

-- Keys
local data_key = KEYS[1]

-- Check if this is a non-refilling bucket
local is_non_refilling = (refill_amount == 0 or refill_frequency == 0)

-- Get current time in milliseconds from Redis
local time_parts = redis.call('TIME')
local now = tonumber(time_parts[1]) * 1000 + (tonumber(time_parts[2]) / 1000)

-- Default bucket values (used if no bucket exists yet)
local tokens = math.min(initial_tokens, capacity)
local slot = now

-- If window_start_time is set, align the initial slot
if window_start_time > 0 then
    local slots_since_start = math.floor((now - window_start_time) / refill_frequency)
    slot = window_start_time + slots_since_start * refill_frequency
end

-- Retrieve stored state, if any
local data = redis.call('GET', data_key)
if data then
    local last_slot, stored_tokens = data:match('(%S+) (%S+)')
    slot = tonumber(last_slot)
    tokens = tonumber(stored_tokens)

    -- Refill tokens based on elapsed time (skip for non-refilling buckets)
    if not is_non_refilling then
        -- Calculate the number of slots that have passed since the last update
        local slots_passed = math.floor((now - slot) / refill_frequency)
        if slots_passed > 0 then
            -- Refill the tokens based on the number of slots passed, capped by capacity
            tokens = math.min(tokens + slots_passed * refill_amount, capacity)
            -- Update the slot to the current aligned slot
            if window_start_time > 0 then
                local slots_since_start = math.floor((now - window_start_time) / refill_frequency)
                slot = window_start_time + slots_since_start * refill_frequency
            else
                slot = now
            end
        end
    end
end

-- If not enough tokens are available, move to the next slot(s) and refill accordingly
if tokens < tokens_to_consume then
    if is_non_refilling then
        -- Non-refilling bucket has run out of tokens
        return redis.error_reply(string.format(
            "No tokens available. Available: %.2f, Requested: %.2f. This is a non-refilling bucket.",
            tokens, tokens_to_consume
        ))
    end
    -- Calculate how many additional tokens we need
    local needed_tokens = tokens_to_consume - tokens
    -- Calculate how many slots we need to move forward to get enough tokens
    local needed_slots = math.ceil(needed_tokens / refill_amount)
    slot = slot + needed_slots * refill_frequency
    tokens = tokens + needed_slots * refill_amount
    -- Clamp tokens to capacity
    tokens = math.min(tokens, capacity)
end

-- Validate max_sleep BEFORE consuming tokens
local required_sleep = math.max(0, slot - now)
-- Check if sleep would exceed max_sleep (if max_sleep > 0)
if max_sleep_ms > 0 and required_sleep > max_sleep_ms then
    return redis.error_reply("Time till next token exceeds max_sleep time:" .. string.format("%.2f", required_sleep/1000)) -- Convert to seconds
end

-- Consume tokens
tokens = tokens - tokens_to_consume

-- Save updated state and set expiry
redis.call('SETEX', data_key, expiry, string.format('%.4f %.4f', slot, tokens))

-- Return the slot when the next token(s) will be available
return slot
