local operation = ARGV[1]
local max_safe_integer = 9007199254740991

local function now_ms()
    local value = redis.call('TIME')
    return tonumber(value[1]) * 1000 + math.floor(tonumber(value[2]) / 1000)
end

local function reply(status, generation, epoch, wait_ms)
    return cjson.encode({
        status = status,
        generation = generation,
        epoch = tostring(epoch or '0'),
        waitMs = wait_ms
    })
end

local function read_leader()
    local encoded = redis.call('GET', KEYS[3])
    if not encoded then
        return nil, nil
    end
    local ok, value = pcall(cjson.decode, encoded)
    if not ok or type(value) ~= 'table' or type(value.generation) ~= 'string' or
        type(value.holderId) ~= 'string' or type(value.token) ~= 'string' or
        type(value.epoch) ~= 'string' then
        return nil, 'invalid'
    end
    return value, nil
end

local function same_leader(value, generation, holder_id, token, epoch)
    return value and value.generation == generation and value.holderId == holder_id and
        value.token == token and value.epoch == tostring(epoch)
end

local function read_counter(key)
    local encoded = redis.call('GET', key)
    if not encoded then
        return 0, nil
    end
    local value = tonumber(encoded)
    if not value or value < 0 or value ~= math.floor(value) then
        return nil, 'invalid'
    end
    return value, nil
end

local function read_permit(token)
    local encoded = redis.call('HGET', KEYS[8], token)
    if not encoded then
        return nil, nil
    end
    local ok, value = pcall(cjson.decode, encoded)
    if not ok or type(value) ~= 'table' or type(value.state) ~= 'string' or
        type(value.expected) ~= 'number' or value.expected <= 0 or
        value.expected ~= math.floor(value.expected) then
        return nil, 'invalid'
    end
    return value, nil
end

local function cleanup_expired(now)
    local expired = redis.call('ZRANGEBYSCORE', KEYS[9], '-inf', now, 'LIMIT', 0, 256)
    if #expired == 0 then
        return nil
    end
    local fetched, fetched_error = read_counter(KEYS[6])
    local pending, pending_error = read_counter(KEYS[7])
    if fetched_error or pending_error then
        return 'invalid'
    end
    for _, permit_token in ipairs(expired) do
        local permit, permit_error = read_permit(permit_token)
        if permit_error then
            return 'invalid'
        end
        if permit and permit.state == 'reserved' then
            if pending < permit.expected or fetched > max_safe_integer - permit.expected then
                return 'invalid'
            end
            pending = pending - permit.expected
            fetched = fetched + permit.expected
        end
        redis.call('HDEL', KEYS[8], permit_token)
        redis.call('ZREM', KEYS[5], permit_token)
        redis.call('ZREM', KEYS[9], permit_token)
    end
    redis.call('SET', KEYS[6], fetched)
    redis.call('SET', KEYS[7], pending)
    return nil
end

if operation == 'acquire_leader' then
    local generation_candidate = ARGV[2]
    local holder_id = ARGV[3]
    local token = ARGV[4]
    local lease_ttl = tonumber(ARGV[5])
    if not lease_ttl or lease_ttl <= 0 then
        return reply('invalid', '', 0, 0)
    end
    local created = redis.call('SET', KEYS[1], generation_candidate, 'NX')
    if created then
        redis.call('DEL', KEYS[2], KEYS[3], KEYS[4], KEYS[5], KEYS[6], KEYS[7], KEYS[8], KEYS[9])
    end
    local generation = redis.call('GET', KEYS[1])
    local current, current_error = read_leader()
    if current_error then
        return reply('unavailable', generation, 0, 0)
    end
    if current then
        if current.generation == generation and current.holderId == holder_id and current.token == token then
            current.expiresAtMs = now_ms() + lease_ttl
            redis.call('SET', KEYS[3], cjson.encode(current), 'PX', lease_ttl)
            return reply('ok', generation, current.epoch, 0)
        end
        local remaining = redis.call('PTTL', KEYS[3])
        return reply('busy', generation, 0, math.max(remaining, 1))
    end
    redis.call('INCR', KEYS[2])
    local epoch = redis.call('GET', KEYS[2])
    local leader = {
        version = 1,
        generation = generation,
        holderId = holder_id,
        token = token,
        epoch = epoch,
        expiresAtMs = now_ms() + lease_ttl
    }
    redis.call('SET', KEYS[3], cjson.encode(leader), 'PX', lease_ttl)
    return reply('ok', generation, epoch, 0)
end

if operation == 'renew_leader' then
    local generation = ARGV[2]
    local holder_id = ARGV[3]
    local token = ARGV[4]
    local epoch = ARGV[5]
    local lease_ttl = tonumber(ARGV[6])
    local current, current_error = read_leader()
    if current_error then
        return reply('unavailable', generation, epoch or 0, 0)
    end
    if not same_leader(current, generation, holder_id, token, epoch) then
        return reply('stale', generation, epoch or 0, 0)
    end
    current.expiresAtMs = now_ms() + lease_ttl
    redis.call('SET', KEYS[3], cjson.encode(current), 'PX', lease_ttl)
    return reply('ok', generation, epoch, 0)
end

if operation == 'release_leader' then
    local generation = ARGV[2]
    local holder_id = ARGV[3]
    local token = ARGV[4]
    local epoch = ARGV[5]
    local current, current_error = read_leader()
    if current_error then
        return reply('unavailable', generation, epoch or 0, 0)
    end
    if same_leader(current, generation, holder_id, token, epoch) then
        redis.call('DEL', KEYS[3])
    end
    return reply('ok', generation, epoch or 0, 0)
end

if operation == 'acquire_fetch' then
    local generation = ARGV[2]
    local holder_id = ARGV[3]
    local leader_token = ARGV[4]
    local epoch = ARGV[5]
    local permit_token = ARGV[6]
    local request_interval = tonumber(ARGV[7])
    local max_in_flight = tonumber(ARGV[8])
    local permit_ttl = tonumber(ARGV[9])
    local expected = tonumber(ARGV[10])
    local max_total = tonumber(ARGV[11])
    if not expected or expected <= 0 or expected ~= math.floor(expected) or
        expected > max_safe_integer or not max_total or max_total < 0 or
        max_total > max_safe_integer or max_total ~= math.floor(max_total) then
        return reply('invalid', generation, epoch or 0, 0)
    end
    local current, current_error = read_leader()
    if current_error then
        return reply('unavailable', generation, epoch or 0, 0)
    end
    if not same_leader(current, generation, holder_id, leader_token, epoch) then
        return reply('stale', generation, epoch or 0, 0)
    end
    local now = now_ms()
    if cleanup_expired(now) then
        return reply('unavailable', generation, epoch, 0)
    end
    local existing, existing_error = read_permit(permit_token)
    if existing_error then
        return reply('unavailable', generation, epoch, 0)
    end
    if existing and existing.state == 'reserved' then
        return reply('ok', generation, epoch, 0)
    end
    if existing then
        return reply('unavailable', generation, epoch, 0)
    end
    local fetched, fetched_error = read_counter(KEYS[6])
    local pending, pending_error = read_counter(KEYS[7])
    if fetched_error or pending_error then
        return reply('unavailable', generation, epoch, 0)
    end
    if max_total > 0 and fetched + pending + expected > max_total then
        return reply('quota_exhausted', generation, epoch, 0)
    end
    local next_request = redis.call('GET', KEYS[4])
    if next_request and not tonumber(next_request) then
        return reply('unavailable', generation, epoch, 0)
    end
    if next_request and tonumber(next_request) > now then
        return reply('rate_limited', generation, epoch, tonumber(next_request) - now)
    end
    if redis.call('ZCARD', KEYS[5]) >= max_in_flight then
        local earliest = redis.call('ZRANGE', KEYS[5], 0, 0, 'WITHSCORES')
        local wait_ms = 1
        if earliest[2] then
            wait_ms = math.max(tonumber(earliest[2]) - now, 1)
        end
        return reply('at_capacity', generation, epoch, wait_ms)
    end
    redis.call('ZADD', KEYS[5], now + permit_ttl, permit_token)
    redis.call('ZADD', KEYS[9], now + permit_ttl, permit_token)
    redis.call('HSET', KEYS[8], permit_token, cjson.encode({
        version = 1,
        state = 'reserved',
        expected = expected
    }))
    redis.call('SET', KEYS[7], pending + expected)
    if request_interval > 0 then
        redis.call('SET', KEYS[4], now + request_interval, 'PX', request_interval)
    else
        redis.call('DEL', KEYS[4])
    end
    return reply('ok', generation, epoch, 0)
end

if operation == 'complete_fetch' or operation == 'cancel_fetch' then
    local permit_token = ARGV[2]
    local fetched_count = tonumber(ARGV[3])
    local settlement_ttl = tonumber(ARGV[4])
    if not fetched_count or fetched_count < 0 or fetched_count ~= math.floor(fetched_count) or
        fetched_count > max_safe_integer or not settlement_ttl or settlement_ttl <= 0 then
        return reply('invalid', '', 0, 0)
    end
    local now = now_ms()
    if cleanup_expired(now) then
        return reply('unavailable', '', 0, 0)
    end
    local permit, permit_error = read_permit(permit_token)
    if permit_error then
        return reply('unavailable', '', 0, 0)
    end
    if not permit or permit.state ~= 'reserved' then
        return reply('ok', '', 0, 0)
    end
    local fetched, fetched_error = read_counter(KEYS[6])
    local pending, pending_error = read_counter(KEYS[7])
    if fetched_error or pending_error or pending < permit.expected or
        fetched > max_safe_integer - fetched_count then
        return reply('unavailable', '', 0, 0)
    end
    pending = pending - permit.expected
    if operation == 'complete_fetch' then
        fetched = fetched + fetched_count
        permit.state = 'completed'
    else
        permit.state = 'cancelled'
    end
    redis.call('SET', KEYS[6], fetched)
    redis.call('SET', KEYS[7], pending)
    redis.call('HSET', KEYS[8], permit_token, cjson.encode(permit))
    redis.call('ZREM', KEYS[5], permit_token)
    redis.call('ZADD', KEYS[9], now + settlement_ttl, permit_token)
    return reply('ok', '', 0, 0)
end

return reply('invalid', '', 0, 0)