local records_key = KEYS[1]
local inventory_key = KEYS[2]
local available_upstream_key = KEYS[3]
local owned_upstream_key = KEYS[4]
local owners_key = KEYS[5]
local sessions_key = KEYS[6]
local session_expiry_key = KEYS[7]
local runtime_key = KEYS[8]
local runtime_expiry_key = KEYS[9]

local upstream_id = ARGV[1]
local safety_margin_ms = tonumber(ARGV[2])
local scan_limit = tonumber(ARGV[3])
local cleanup_limit = tonumber(ARGV[4])

local function reply(status, managed, available_slots)
    return cjson.encode({
        status = status,
        managed = managed or 0,
        availableSlots = tostring(available_slots or 0)
    })
end

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

local function decode_table(value)
    if not value then
        return nil
    end
    local ok, decoded = pcall(cjson.decode, value)
    if not ok or type(decoded) ~= 'table' then
        return nil
    end
    return decoded
end

local function valid_uint(value)
    return type(value) == 'string' and string.match(value, '^[0-9]+$') and
        value ~= '0' and (string.len(value) == 1 or string.sub(value, 1, 1) ~= '0')
end

local function compare_uint(left, right)
    if string.len(left) ~= string.len(right) then
        return string.len(left) < string.len(right) and -1 or 1
    end
    if left == right then
        return 0
    end
    return left < right and -1 or 1
end

local function cleanup(now)
    local expired_sessions = redis.call('ZRANGEBYSCORE', session_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit)
    for _, worker_id in ipairs(expired_sessions) do
        redis.call('HDEL', sessions_key, worker_id)
        redis.call('HDEL', runtime_key, worker_id)
        redis.call('ZREM', session_expiry_key, worker_id)
        redis.call('ZREM', runtime_expiry_key, worker_id)
    end
    local expired_reports = redis.call('ZRANGEBYSCORE', runtime_expiry_key, '-inf', now, 'LIMIT', 0, cleanup_limit)
    for _, worker_id in ipairs(expired_reports) do
        redis.call('HDEL', runtime_key, worker_id)
        redis.call('ZREM', runtime_expiry_key, worker_id)
    end
end

if type(upstream_id) ~= 'string' or upstream_id == '' or not safety_margin_ms or safety_margin_ms < 0 or
    not scan_limit or scan_limit <= 0 or not cleanup_limit or cleanup_limit <= 0 then
    return reply('invalid', 0, 0)
end

local now = now_ms()
cleanup(now)
local threshold = now + safety_margin_ms
local available_ids = redis.call('ZRANGEBYSCORE', available_upstream_key, '(' .. threshold, '+inf',
    'LIMIT', 0, scan_limit + 1)
if #available_ids > scan_limit then
    return reply('unavailable', 0, 0)
end
local remaining = scan_limit - #available_ids
local owned_ids = redis.call('ZRANGEBYSCORE', owned_upstream_key, '(' .. threshold, '+inf',
    'LIMIT', 0, remaining + 1)
if #owned_ids > remaining then
    return reply('unavailable', 0, 0)
end

local managed = tonumber(redis.call('HGET', inventory_key, upstream_id) or '0')
if not managed or managed < 0 or managed ~= math.floor(managed) then
    return reply('unavailable', 0, 0)
end
local available_slots = 0
local worker_cache = {}
local seen = {}
for _, proxy_id in ipairs(available_ids) do
    seen[proxy_id] = true
    local record = decode_table(redis.call('HGET', records_key, proxy_id))
    if not record or type(record.sourceUpstream) ~= 'string' or type(record.state) ~= 'string' or
        record.sourceUpstream ~= upstream_id or record.state ~= 'AVAILABLE' or
        type(record.usableUntilMs) ~= 'number' or record.usableUntilMs <= threshold or
        type(record.maxConcurrency) ~= 'number' or record.maxConcurrency < 0 or
        record.maxConcurrency ~= math.floor(record.maxConcurrency) or
        (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HGET', owners_key, proxy_id) then
        return reply('unavailable', 0, 0)
    end
    available_slots = available_slots + record.maxConcurrency
end

for _, proxy_id in ipairs(owned_ids) do
    if seen[proxy_id] then
        return reply('unavailable', 0, 0)
    end
    local record = decode_table(redis.call('HGET', records_key, proxy_id))
    if not record or record.sourceUpstream ~= upstream_id or record.state ~= 'AVAILABLE' or
        type(record.usableUntilMs) ~= 'number' or record.usableUntilMs <= threshold or
        type(record.maxConcurrency) ~= 'number' or record.maxConcurrency < 0 or
        record.maxConcurrency ~= math.floor(record.maxConcurrency) or
        type(record.ownerWorkerId) ~= 'string' or record.ownerWorkerId == '' then
        return reply('unavailable', 0, 0)
    end
    local owner_worker_id = record.ownerWorkerId
    local owner = decode_table(redis.call('HGET', owners_key, proxy_id))
    if not owner or owner.workerId ~= owner_worker_id or type(owner.epoch) ~= 'number' or
        type(owner.expiresAtMs) ~= 'number' or owner.expiresAtMs <= now or
        type(owner.draining) ~= 'boolean' then
        return reply('unavailable', 0, 0)
    end
    local cached = worker_cache[owner_worker_id]
    if not cached then
        local session = decode_table(redis.call('HGET', sessions_key, owner_worker_id))
        local report = decode_table(redis.call('HGET', runtime_key, owner_worker_id))
        cached = {fresh = false, counters = {}}
        if session and report and session.workerId == owner_worker_id and
            report.workerId == owner_worker_id and session.sessionId == report.sessionId and
            type(session.expiresAtMs) == 'number' and session.expiresAtMs > now and
            type(report.expiresAtMs) == 'number' and report.expiresAtMs > now and
            valid_uint(session.ackedSnapshotVersion) and valid_uint(session.ackedOwnershipEpoch) and
            report.snapshotVersion == session.ackedSnapshotVersion and
            report.ownershipEpoch == session.ackedOwnershipEpoch then
            cached.fresh = true
            cached.ownershipEpoch = report.ownershipEpoch
            if type(report.counters) == 'table' then
                for _, counter in pairs(report.counters) do
                    if type(counter) == 'table' and type(counter.proxyId) == 'string' then
                        cached.counters[counter.proxyId] = counter
                    end
                end
            end
        end
        worker_cache[owner_worker_id] = cached
    end
    local owner_epoch = tostring(owner.epoch)
    if cached.fresh and valid_uint(owner_epoch) and
        compare_uint(cached.ownershipEpoch, owner_epoch) >= 0 then
        local counter = cached.counters[proxy_id]
        local active = 0
        local reserved = 0
        local draining = false
        if counter then
            active = counter.active
            reserved = counter.reserved
            draining = counter.draining
        end
        if type(active) == 'number' and type(reserved) == 'number' and active >= 0 and reserved >= 0 and
            active == math.floor(active) and reserved == math.floor(reserved) and
            not draining and not owner.draining then
            local slots = record.maxConcurrency - active - reserved
            if slots > 0 then
                available_slots = available_slots + slots
            end
        end
    end
end

return reply('ok', managed, available_slots)