local operation = ARGV[1]

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

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
    redis.call('SET', KEYS[1], generation_candidate, 'NX')
    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 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()
    redis.call('ZREMRANGEBYSCORE', KEYS[5], '-inf', now)
    local existing = redis.call('ZSCORE', KEYS[5], permit_token)
    if existing then
        return reply('ok', 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('PEXPIRE', KEYS[5], permit_ttl + 1000)
    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 == 'release_fetch' then
    redis.call('ZREM', KEYS[5], ARGV[2])
    return reply('ok', '', 0, 0)
end

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