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 epoch_key = KEYS[10]
local operation_key = KEYS[11]

local operation = ARGV[1]
local operation_ttl_ms = tonumber(ARGV[2])
local cleanup_limit = tonumber(ARGV[3])
local now_ms = tonumber(ARGV[4])
local proxy_id = ARGV[5]
local worker_id = ARGV[6]
local epoch = tonumber(ARGV[7])
local value = tonumber(ARGV[8])
local active = tonumber(ARGV[9])
local reserved = tonumber(ARGV[10])
local mutating = operation ~= 'get'

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

if mutating then
    local committed = redis.call('GET', operation_key)
    if committed then
        return committed
    end
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 type(upstream) ~= 'string' or upstream == '' then
        return
    end
    local count = redis.call('HINCRBY', inventory_key, upstream, -1)
    if count < 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 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 count = redis.call('HINCRBY', state_inventory_key, field, -1)
    if count <= 0 then
        redis.call('HDEL', state_inventory_key, field)
    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 remove_available(id, record)
    redis.call('ZREM', available_key, id)
    local index_keys = record and record.indexKeys
    if type(index_keys) == 'table' then
        for _, index_key in ipairs(index_keys) do
            if type(index_key) == 'string' and index_key ~= '' then
                redis.call('ZREM', index_key, id)
            end
        end
    end
end

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

local function add_available(id, record, at_ms)
    local usable_until_ms = record and tonumber(record.usableUntilMs)
    if not usable_until_ms or record.state ~= 'AVAILABLE' or usable_until_ms <= at_ms then
        return
    end
    redis.call('ZADD', available_key, usable_until_ms, id)
    touch(available_key, tonumber(record.expiresAtMs))
    if type(record.indexKeys) == 'table' then
        for _, index_key in ipairs(record.indexKeys) do
            if type(index_key) == 'string' and index_key ~= '' then
                redis.call('ZADD', index_key, usable_until_ms, id)
                touch(index_key, tonumber(record.expiresAtMs))
            end
        end
    end
end

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

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

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

local function valid_assignment(assignment)
    return assignment and assignment.version == 1 and type(assignment.proxyId) == 'string' and
        assignment.proxyId ~= '' and type(assignment.workerId) == 'string' and assignment.workerId ~= '' and
        tonumber(assignment.epoch) and tonumber(assignment.epoch) > 0 and
        tonumber(assignment.assignmentVersion) and tonumber(assignment.assignmentVersion) > 0 and
        tonumber(assignment.expiresAtMs) and tonumber(assignment.expiresAtMs) > 0 and
        type(assignment.draining) == 'boolean'
end

local function clear_owner(id, assignment, at_ms, restore)
    local raw_record = redis.call('HGET', records_key, id)
    local record = decode_table(raw_record)
    if record and (not assignment or record.ownerWorkerId == assignment.workerId) then
        remove_owned(id, record)
        record.ownerWorkerId = nil
        redis.call('HSET', records_key, id, cjson.encode(record))
        if restore then
            add_available(id, record, at_ms)
        end
    end
    redis.call('HDEL', owners_key, id)
    redis.call('ZREM', owner_expiry_key, id)
end

if operation == 'assign' then
    cleanup_hard_expired(now_ms)
    local current_raw = redis.call('HGET', owners_key, proxy_id)
    if current_raw then
        local current = decode_table(current_raw)
        if not valid_assignment(current) then
            return finish({status = 'unavailable'})
        end
        if tonumber(current.expiresAtMs) > now_ms then
            return finish({status = 'already_owned'})
        end
        clear_owner(proxy_id, current, now_ms, true)
    end

    local record = decode_table(redis.call('HGET', records_key, proxy_id))
    if not record or record.state ~= 'AVAILABLE' or
        type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' or
        (record.ownerWorkerId and record.ownerWorkerId ~= '') or
        redis.call('HEXISTS', owners_key, proxy_id) == 1 or
        not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
        return finish({status = 'unavailable'})
    end
    local expires_at_ms = now_ms + value
    if tonumber(record.usableUntilMs) < expires_at_ms then
        expires_at_ms = tonumber(record.usableUntilMs)
    end
    local next_epoch = redis.call('INCR', epoch_key)
    local assignment = {
        version = 1,
        proxyId = proxy_id,
        workerId = worker_id,
        epoch = next_epoch,
        assignmentVersion = 1,
        expiresAtMs = expires_at_ms,
        draining = false,
    }
    local encoded = cjson.encode(assignment)
    redis.call('HSET', owners_key, proxy_id, encoded)
    redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
    touch(owners_key, tonumber(record.expiresAtMs))
    touch(owner_expiry_key, tonumber(record.expiresAtMs))
    touch(epoch_key, tonumber(record.expiresAtMs))
    record.ownerWorkerId = worker_id
    redis.call('HSET', records_key, proxy_id, cjson.encode(record))
    remove_available(proxy_id, record)
    redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
    touch(record.ownerIndexKey, tonumber(record.expiresAtMs))
    return finish({status = 'ok', record = encoded})
end

if operation == 'renew' then
    cleanup_hard_expired(now_ms)
    local current = decode_table(redis.call('HGET', owners_key, proxy_id))
    if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
        return finish({status = 'stale'})
    end
    if tonumber(current.expiresAtMs) <= now_ms then
        clear_owner(proxy_id, current, now_ms, true)
        return finish({status = 'stale'})
    end
    local record = decode_table(redis.call('HGET', records_key, proxy_id))
    if not record or record.ownerWorkerId ~= worker_id or
        not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
        clear_owner(proxy_id, current, now_ms, false)
        return finish({status = 'stale'})
    end
    local expires_at_ms = now_ms + value
    if tonumber(record.usableUntilMs) < expires_at_ms then
        expires_at_ms = tonumber(record.usableUntilMs)
    end
    current.assignmentVersion = tonumber(current.assignmentVersion) + 1
    current.expiresAtMs = expires_at_ms
    local encoded = cjson.encode(current)
    redis.call('HSET', owners_key, proxy_id, encoded)
    redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
    touch(owners_key, tonumber(record.expiresAtMs))
    touch(owner_expiry_key, tonumber(record.expiresAtMs))
    touch(epoch_key, tonumber(record.expiresAtMs))
    return finish({status = 'ok', record = encoded})
end

if operation == 'begin_drain' then
    local current = decode_table(redis.call('HGET', owners_key, proxy_id))
    if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
        return finish({status = 'stale'})
    end
    if not current.draining then
        current.draining = true
        current.assignmentVersion = tonumber(current.assignmentVersion) + 1
        local encoded = cjson.encode(current)
        redis.call('HSET', owners_key, proxy_id, encoded)
        local record = decode_table(redis.call('HGET', records_key, proxy_id))
        remove_owned(proxy_id, record)
        return finish({status = 'ok', record = encoded})
    end
    return finish({status = 'ok', record = cjson.encode(current)})
end

if operation == 'acknowledge_drain' then
    local current = decode_table(redis.call('HGET', owners_key, proxy_id))
    if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
        return finish({status = 'stale'})
    end
    if not current.draining then
        return finish({status = 'not_draining'})
    end
    if active > 0 or reserved > 0 then
        return finish({status = 'drain_not_ready'})
    end
    local server_time = redis.call('TIME')
    local server_now_ms = tonumber(server_time[1]) * 1000 + math.floor(tonumber(server_time[2]) / 1000)
    clear_owner(proxy_id, current, server_now_ms, true)
    return finish({status = 'ok'})
end

if operation == 'get' then
    local raw = redis.call('HGET', owners_key, proxy_id)
    if not raw then
        return finish({status = 'not_found'})
    end
    return finish({status = 'ok', record = raw})
end

if operation == 'expire' then
    local ids = redis.call('ZRANGEBYSCORE', owner_expiry_key, '-inf', now_ms, 'LIMIT', 0, value)
    local expired = {}
    for _, id in ipairs(ids) do
        local current = decode_table(redis.call('HGET', owners_key, id))
        if valid_assignment(current) then
            expired[#expired + 1] = current
        end
        clear_owner(id, current, now_ms, true)
    end
    local encoded = '[]'
    if #expired > 0 then
        encoded = cjson.encode(expired)
    end
    return finish({status = 'ok', record = encoded})
end

return finish({status = 'invalid'})