local records_key = KEYS[1]
local unique_key = KEYS[2]
local idkeys_key = KEYS[3]
local expiry_key = KEYS[4]
local available_key = KEYS[5]
local inventory_key = KEYS[6]
local state_inventory_key = KEYS[7]
local owners_key = KEYS[8]
local owner_expiry_key = KEYS[9]
local operation_key = KEYS[10]

local now_ms = tonumber(ARGV[1])
local cleanup_limit = tonumber(ARGV[2])
local max_size = tonumber(ARGV[3])
local operation_ttl_ms = tonumber(ARGV[4])
local candidates = cjson.decode(ARGV[5])

local committed = redis.call('GET', operation_key)
if committed then
    return committed
end

local function is_managed(state)
    return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
        state == 'SUSPECT' or state == 'DRAINING'
end

local function decrement_inventory(upstream)
    if not upstream or upstream == '' then
        return
    end
    local value = redis.call('HINCRBY', inventory_key, upstream, -1)
    if value < 0 then
        redis.call('HSET', inventory_key, upstream, 0)
    end
end

local function state_field(upstream, state)
    return string.len(upstream) .. ':' .. upstream .. ':' .. state
end

local function is_counted(state)
    return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
        state == 'SUSPECT' or state == 'DRAINING' or state == 'UNHEALTHY' or state == 'EXTRACTED'
end

local function increment_state(upstream, state)
    if not is_counted(state) then
        return
    end
    redis.call('HINCRBY', state_inventory_key, state_field(upstream, state), 1)
end

local function decrement_state(upstream, state)
    if type(upstream) ~= 'string' or upstream == '' or not is_counted(state) then
        return
    end
    local field = state_field(upstream, state)
    local value = redis.call('HINCRBY', state_inventory_key, field, -1)
    if value <= 0 then
        redis.call('HDEL', state_inventory_key, field)
    end
end

local function remove_available(proxy_id, record)
    redis.call('ZREM', available_key, proxy_id)
    local indexes = record and record.indexKeys or {}
    for _, index_key in ipairs(indexes) do
        redis.call('ZREM', index_key, proxy_id)
    end
end

local function remove_owned(proxy_id, record)
    if record and type(record.ownerIndexKey) == 'string' and record.ownerIndexKey ~= '' then
        redis.call('ZREM', record.ownerIndexKey, proxy_id)
    end
end

local function remove_worker_owned(proxy_id)
    local raw = redis.call('HGET', owners_key, proxy_id)
    if not raw then
        return
    end
    local decoded, assignment = pcall(cjson.decode, raw)
    if decoded and type(assignment) == 'table' and
        type(assignment.workerIndexKey) == 'string' and assignment.workerIndexKey ~= '' then
        redis.call('ZREM', assignment.workerIndexKey, proxy_id)
    end
end

local function sync_worker_owned(proxy_id, record)
    local raw = redis.call('HGET', owners_key, proxy_id)
    if not raw then
        return
    end
    local decoded, assignment = pcall(cjson.decode, raw)
    if not decoded or type(assignment) ~= 'table' or
        type(assignment.workerIndexKey) ~= 'string' or assignment.workerIndexKey == '' then
        return
    end
    if assignment.workerId == record.ownerWorkerId and not assignment.draining and
        type(assignment.expiresAtMs) == 'number' and assignment.expiresAtMs > now_ms and
        record.state == 'AVAILABLE' and tonumber(record.usableUntilMs) > now_ms then
        redis.call('ZADD', assignment.workerIndexKey, assignment.expiresAtMs, proxy_id)
        local current = redis.call('PEXPIRETIME', assignment.workerIndexKey)
        if current < tonumber(record.expiresAtMs) then
            redis.call('PEXPIREAT', assignment.workerIndexKey, tonumber(record.expiresAtMs))
        end
        return
    end
    redis.call('ZREM', assignment.workerIndexKey, proxy_id)
end

local function remove_proxy(proxy_id)
    local raw = redis.call('HGET', records_key, proxy_id)
    local record = nil
    if raw then
        record = cjson.decode(raw)
        remove_available(proxy_id, record)
        remove_owned(proxy_id, record)
        if is_managed(record.state) then
            decrement_inventory(record.sourceUpstream)
        end
        decrement_state(record.sourceUpstream, record.state)
    else
        redis.call('ZREM', available_key, proxy_id)
    end
    local digest = redis.call('HGET', idkeys_key, proxy_id)
    if digest and redis.call('HGET', unique_key, digest) == proxy_id then
        redis.call('HDEL', unique_key, digest)
    end
    redis.call('HDEL', idkeys_key, proxy_id)
    redis.call('HDEL', records_key, proxy_id)
    redis.call('ZREM', expiry_key, proxy_id)
    remove_worker_owned(proxy_id)
    redis.call('HDEL', owners_key, proxy_id)
    redis.call('ZREM', owner_expiry_key, proxy_id)
end

local function cleanup_expired()
    local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', now_ms, 'LIMIT', 0, cleanup_limit)
    for _, proxy_id in ipairs(expired) do
        remove_proxy(proxy_id)
    end
end

local function touch(key, expires_at_ms)
    if redis.call('EXISTS', key) == 0 then
        return
    end
    local current = redis.call('PEXPIRETIME', key)
    if current < expires_at_ms then
        redis.call('PEXPIREAT', key, expires_at_ms)
    end
end

local function add_available(proxy_id, record)
    if record.state ~= 'AVAILABLE' or (record.ownerWorkerId and record.ownerWorkerId ~= '') or
        tonumber(record.usableUntilMs) <= now_ms then
        return
    end
    redis.call('ZADD', available_key, record.usableUntilMs, proxy_id)
    for _, index_key in ipairs(record.indexKeys or {}) do
        redis.call('ZADD', index_key, record.usableUntilMs, proxy_id)
        touch(index_key, tonumber(record.expiresAtMs))
    end
end

local function sync_owned(proxy_id, record)
    if record.state == 'AVAILABLE' and record.ownerWorkerId and record.ownerWorkerId ~= '' and
        tonumber(record.usableUntilMs) > now_ms then
        redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
        touch(record.ownerIndexKey, tonumber(record.expiresAtMs))
    else
        redis.call('ZREM', record.ownerIndexKey, proxy_id)
    end
end

local function finish(reply)
    local encoded = cjson.encode(reply)
    redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
    return encoded
end

cleanup_expired()

for _, candidate in ipairs(candidates) do
    local decoded, incoming = pcall(cjson.decode, candidate.record)
    if not decoded or type(incoming) ~= 'table' or
        type(incoming.ownerIndexKey) ~= 'string' or incoming.ownerIndexKey == '' then
        return finish({status = 'invalid', accepted = 0, inserted = 0, refreshed = 0, dropped = 0})
    end
    local mapped = redis.call('HGET', idkeys_key, candidate.proxyId)
    if mapped and mapped ~= candidate.uniqueDigest then
        return finish({status = 'invalid', accepted = 0, inserted = 0, refreshed = 0, dropped = 0})
    end
end

local accepted = #candidates
local inserted = 0
local refreshed = 0
local dropped = 0
local max_expiry_ms = 0

for _, candidate in ipairs(candidates) do
    local incumbent_id = redis.call('HGET', unique_key, candidate.uniqueDigest)
    local current_raw = incumbent_id and redis.call('HGET', records_key, incumbent_id) or nil
    if incumbent_id and not current_raw then
        remove_proxy(incumbent_id)
        redis.call('HDEL', unique_key, candidate.uniqueDigest)
        incumbent_id = nil
    end

    if current_raw then
        local current = cjson.decode(current_raw)
        if tonumber(current.expiresAtMs) <= now_ms then
            remove_proxy(incumbent_id)
            incumbent_id = nil
            current_raw = nil
        elseif current.state == 'EXTRACTED' or current.sourceUpstream ~= candidate.upstream then
            refreshed = refreshed + 1
        else
            local incoming = cjson.decode(candidate.record)
            remove_available(incumbent_id, current)
            incoming.id = current.id
            incoming.createdAtMs = current.createdAtMs
            incoming.state = current.state
            incoming.lastCheckedAtMs = current.lastCheckedAtMs
            incoming.lastSuccessAtMs = current.lastSuccessAtMs
            incoming.latencyNs = current.latencyNs
            incoming.ownerWorkerId = current.ownerWorkerId
            local encoded = cjson.encode(incoming)
            redis.call('HSET', records_key, incumbent_id, encoded)
            redis.call('ZADD', expiry_key, incoming.expiresAtMs, incumbent_id)
            sync_owned(incumbent_id, incoming)
            sync_worker_owned(incumbent_id, incoming)
            add_available(incumbent_id, incoming)
            if tonumber(incoming.expiresAtMs) > max_expiry_ms then
                max_expiry_ms = tonumber(incoming.expiresAtMs)
            end
            refreshed = refreshed + 1
        end
    end

    if not incumbent_id then
        local current_size = tonumber(redis.call('HGET', inventory_key, candidate.upstream) or '0')
        if current_size >= max_size then
            dropped = dropped + 1
        else
            local incoming = cjson.decode(candidate.record)
            redis.call('HSET', records_key, candidate.proxyId, candidate.record)
            redis.call('HSET', unique_key, candidate.uniqueDigest, candidate.proxyId)
            redis.call('HSET', idkeys_key, candidate.proxyId, candidate.uniqueDigest)
            redis.call('ZADD', expiry_key, incoming.expiresAtMs, candidate.proxyId)
            sync_owned(candidate.proxyId, incoming)
            sync_worker_owned(candidate.proxyId, incoming)
            if is_managed(incoming.state) then
                redis.call('HINCRBY', inventory_key, candidate.upstream, 1)
            end
            increment_state(candidate.upstream, incoming.state)
            add_available(candidate.proxyId, incoming)
            if tonumber(incoming.expiresAtMs) > max_expiry_ms then
                max_expiry_ms = tonumber(incoming.expiresAtMs)
            end
            inserted = inserted + 1
        end
    end
end

if max_expiry_ms > 0 then
    touch(records_key, max_expiry_ms)
    touch(unique_key, max_expiry_ms)
    touch(idkeys_key, max_expiry_ms)
    touch(expiry_key, max_expiry_ms)
    touch(available_key, max_expiry_ms)
    touch(inventory_key, max_expiry_ms)
    touch(state_inventory_key, max_expiry_ms)
    touch(owners_key, max_expiry_ms)
    touch(owner_expiry_key, max_expiry_ms)
end

return finish({
    status = 'ok', accepted = accepted, inserted = inserted,
    refreshed = refreshed, dropped = dropped,
})