local sessions_key = KEYS[1]
local session_expiry_key = KEYS[2]
local runtime_key = KEYS[3]
local runtime_expiry_key = KEYS[4]
local owners_key = KEYS[5]

local operation = ARGV[1]
local ttl_ms = tonumber(ARGV[2])
local cleanup_limit = tonumber(ARGV[3])
local payload = ARGV[4]
local digest = ARGV[5]

local function reply(status, snapshots)
    if snapshots then
        return cjson.encode({status = status, snapshots = snapshots})
    end
    return '{"status":' .. cjson.encode(status) .. ',"snapshots":[]}'
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

local function valid_session(value)
    return value and value.version == 1 and type(value.workerId) == 'string' and value.workerId ~= '' and
        type(value.instanceId) == 'string' and value.instanceId ~= '' and
        type(value.sessionId) == 'string' and value.sessionId ~= '' and
        valid_uint(value.ackedSnapshotVersion) and valid_uint(value.ackedOwnershipEpoch)
end

local function valid_owner(value, worker_id, ownership_epoch, now)
    return value and type(value.workerId) == 'string' and value.workerId == worker_id and
        type(value.epoch) == 'number' and valid_uint(tostring(value.epoch)) and
        compare_uint(ownership_epoch, tostring(value.epoch)) >= 0 and
        type(value.expiresAtMs) == 'number' and value.expiresAtMs > now
end

local now = now_ms()
cleanup(now)

if operation == 'replace_session' then
    if not ttl_ms or ttl_ms <= 0 then
        return reply('invalid')
    end
    local session = decode_table(payload)
    if not valid_session(session) then
        return reply('invalid')
    end
    local current = decode_table(redis.call('HGET', sessions_key, session.workerId))
    if current and valid_session(current) and current.sessionId == session.sessionId and
        current.instanceId == session.instanceId and type(current.expiresAtMs) == 'number' and
        current.expiresAtMs > now then
        local epoch_order = compare_uint(session.ackedOwnershipEpoch, current.ackedOwnershipEpoch)
        local version_order = compare_uint(session.ackedSnapshotVersion, current.ackedSnapshotVersion)
        if epoch_order < 0 or (epoch_order == 0 and version_order < 0) then
            return reply('stale')
        end
        if epoch_order > 0 or version_order > 0 then
            redis.call('HDEL', runtime_key, session.workerId)
            redis.call('ZREM', runtime_expiry_key, session.workerId)
        end
    else
        redis.call('HDEL', runtime_key, session.workerId)
        redis.call('ZREM', runtime_expiry_key, session.workerId)
    end
    session.expiresAtMs = now + ttl_ms
    redis.call('HSET', sessions_key, session.workerId, cjson.encode(session))
    redis.call('ZADD', session_expiry_key, session.expiresAtMs, session.workerId)
    return reply('ok')
end

if operation == 'replace_report' then
    if not ttl_ms or ttl_ms <= 0 or type(digest) ~= 'string' or digest == '' then
        return reply('invalid')
    end
    local report = decode_table(payload)
    if not report or report.version ~= 1 or type(report.workerId) ~= 'string' or report.workerId == '' or
        type(report.sessionId) ~= 'string' or report.sessionId == '' or not valid_uint(report.sequence) or
        not valid_uint(report.snapshotVersion) or not valid_uint(report.ownershipEpoch) or
        type(report.observedAtMs) ~= 'number' or type(report.counters) ~= 'table' then
        return reply('invalid')
    end
    local session = decode_table(redis.call('HGET', sessions_key, report.workerId))
    if not valid_session(session) or session.workerId ~= report.workerId or session.sessionId ~= report.sessionId or
        type(session.expiresAtMs) ~= 'number' or session.expiresAtMs <= now then
        return reply('unavailable')
    end
    if report.snapshotVersion ~= session.ackedSnapshotVersion or
        report.ownershipEpoch ~= session.ackedOwnershipEpoch then
        return reply('stale')
    end
    local current = decode_table(redis.call('HGET', runtime_key, report.workerId))
    if current and current.sessionId == report.sessionId and valid_uint(current.sequence) then
        local ordering = compare_uint(report.sequence, current.sequence)
        if ordering < 0 then
            return reply('stale')
        end
        if ordering == 0 then
            if current.digest == digest then
                return reply('ok')
            end
            return reply('conflict')
        end
    end
    local seen = {}
    for _, counter in pairs(report.counters) do
        if type(counter) ~= 'table' or type(counter.proxyId) ~= 'string' or counter.proxyId == '' or
            type(counter.active) ~= 'number' or counter.active < 0 or counter.active ~= math.floor(counter.active) or
            type(counter.reserved) ~= 'number' or counter.reserved < 0 or counter.reserved ~= math.floor(counter.reserved) or
            type(counter.draining) ~= 'boolean' or seen[counter.proxyId] then
            return reply('invalid')
        end
        seen[counter.proxyId] = true
        local owner = decode_table(redis.call('HGET', owners_key, counter.proxyId))
        if not valid_owner(owner, report.workerId, report.ownershipEpoch, now) then
            return reply('stale')
        end
    end
    report.digest = digest
    report.expiresAtMs = now + ttl_ms
    redis.call('HSET', runtime_key, report.workerId, cjson.encode(report))
    redis.call('ZADD', runtime_expiry_key, report.expiresAtMs, report.workerId)
    session.expiresAtMs = now + ttl_ms
    redis.call('HSET', sessions_key, report.workerId, cjson.encode(session))
    redis.call('ZADD', session_expiry_key, session.expiresAtMs, report.workerId)
    return reply('ok')
end

if operation == 'read' then
    local queries = decode_table(payload)
    if not queries then
        return reply('invalid')
    end
    if next(queries) == nil then
        return '{"status":"ok","snapshots":[]}'
    end
    local snapshots = cjson.decode('[]')
    local cache = {}
    for _, query in ipairs(queries) do
        if type(query) ~= 'table' or type(query.proxyId) ~= 'string' or query.proxyId == '' or
            type(query.workerId) ~= 'string' or query.workerId == '' or not valid_uint(query.ownershipEpoch) then
            return reply('invalid')
        end
        local snapshot = {proxyId = query.proxyId, active = 0, reserved = 0, draining = false, fresh = false}
        local owner = decode_table(redis.call('HGET', owners_key, query.proxyId))
        if valid_owner(owner, query.workerId, query.ownershipEpoch, now) and
            compare_uint(query.ownershipEpoch, tostring(owner.epoch)) == 0 then
            local cached = cache[query.workerId]
            if not cached then
                local session = decode_table(redis.call('HGET', sessions_key, query.workerId))
                local report = decode_table(redis.call('HGET', runtime_key, query.workerId))
                cached = {fresh = false, counters = {}}
                if valid_session(session) and session.workerId == query.workerId and report and
                    report.workerId == query.workerId and report.sessionId == session.sessionId and
                    type(session.expiresAtMs) == 'number' and session.expiresAtMs > now and
                    type(report.expiresAtMs) == 'number' and report.expiresAtMs > now and
                    valid_uint(report.ownershipEpoch) 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
                cache[query.workerId] = cached
            end
            if cached.fresh and compare_uint(cached.ownershipEpoch, query.ownershipEpoch) >= 0 then
                snapshot.fresh = true
                local counter = cached.counters[query.proxyId]
                if counter then
                    snapshot.active = counter.active
                    snapshot.reserved = counter.reserved
                    snapshot.draining = counter.draining
                end
            end
        end
        snapshots[#snapshots + 1] = snapshot
    end
    return reply('ok', snapshots)
end

return reply('invalid')