local records_key = KEYS[1]
local expiry_key = KEYS[2]
local due_key = KEYS[3]
local queued_key = KEYS[4]
local leases_key = KEYS[5]
local tasks_key = KEYS[6]
local task_expiry_key = KEYS[7]
local ref_task_key = KEYS[8]
local checker_leases_key = KEYS[9]
local state_inventory_key = KEYS[10]

local operation = ARGV[1]
local now_ms = tonumber(ARGV[2])
local decoded, payload = pcall(cjson.decode, ARGV[3])
if not decoded or type(payload) ~= 'table' or not now_ms then
    return cjson.encode({status = 'invalid'})
end

local limit = tonumber(payload.limit or 0)
if not limit or limit <= 0 then
    return cjson.encode({status = 'invalid'})
end

local function finish(reply)
    return cjson.encode(reply)
end

local function live_record(proxy_id)
    local raw = redis.call('HGET', records_key, proxy_id)
    if not raw then
        return nil, nil
    end
    local valid, record = pcall(cjson.decode, raw)
    if not valid or type(record) ~= 'table' or tonumber(record.expiresAtMs) <= now_ms or record.state == 'EXTRACTED' then
        return nil, nil
    end
    return raw, record
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_task(task_id, requeue)
    local raw = redis.call('HGET', tasks_key, task_id)
    local task = nil
    if raw then
        local valid
        valid, task = pcall(cjson.decode, raw)
        if not valid or type(task) ~= 'table' then
            task = nil
        end
    end
    redis.call('ZREM', queued_key, task_id)
    redis.call('ZREM', leases_key, task_id)
    redis.call('ZREM', task_expiry_key, task_id)
    redis.call('HDEL', tasks_key, task_id)
    if task then
		if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then
			redis.call('ZREM', task.upstreamTasksKey, task_id)
		end
        if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
            redis.call('ZREM', task.checkerLeaseKey, task_id)
        end
        if type(task.proxyId) == 'string' and redis.call('HGET', ref_task_key, task.proxyId) == task_id then
            redis.call('HDEL', ref_task_key, task.proxyId)
        end
        if requeue and type(task.proxyId) == 'string' then
            local _, record = live_record(task.proxyId)
            if record then
                redis.call('ZADD', due_key, now_ms, task.proxyId)
                touch(due_key, tonumber(record.expiresAtMs))
            end
        end
    end
end

local function reap()
    local expired_leases = redis.call('ZRANGEBYSCORE', leases_key, '-inf', now_ms, 'LIMIT', 0, limit)
    for _, task_id in ipairs(expired_leases) do
        local raw = redis.call('HGET', tasks_key, task_id)
        if not raw then
            redis.call('ZREM', leases_key, task_id)
        else
            local valid, task = pcall(cjson.decode, raw)
            if not valid or type(task) ~= 'table' then
                remove_task(task_id, true)
            elseif task.state == 'LEASED' and tonumber(task.leaseExpiresAtMs or 0) <= now_ms then
                redis.call('ZREM', leases_key, task_id)
                if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
                    redis.call('ZREM', task.checkerLeaseKey, task_id)
                end
                task.state = 'QUEUED'
                task.leaseCheckerId = ''
                task.leaseToken = ''
                task.leaseExpiresAtMs = 0
                task.checkerLeaseKey = ''
                redis.call('HSET', tasks_key, task_id, cjson.encode(task))
                redis.call('ZADD', queued_key, tonumber(task.priority), task_id)
            else
                redis.call('ZREM', leases_key, task_id)
            end
        end
    end
    local expired_tasks = redis.call('ZRANGEBYSCORE', task_expiry_key, '-inf', now_ms, 'LIMIT', 0, limit)
    for _, task_id in ipairs(expired_tasks) do
        local raw = redis.call('HGET', tasks_key, task_id)
        local requeue = true
        if raw then
            local valid, task = pcall(cjson.decode, raw)
            if valid and type(task) == 'table' and task.state == 'DONE' then
                requeue = false
            end
        end
        remove_task(task_id, requeue)
    end
end

local function valid_state(state)
    return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
        state == 'SUSPECT' or state == 'UNHEALTHY'
end

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

local function decrement_state(upstream, state)
    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 increment_state(upstream, state)
    redis.call('HINCRBY', state_inventory_key, state_field(upstream, state), 1)
end

local function matches_task(task, fact)
    return task and task.taskId == fact.taskId and task.proxyId == fact.proxyId and task.level == fact.level and
        (task.routingName or '') == (fact.routingName or '') and (task.targetUrl or '') == (fact.targetUrl or '')
end

local function authorize(fact)
    if type(fact) ~= 'table' or type(fact.taskId) ~= 'string' or type(fact.proxyId) ~= 'string' or
        type(fact.level) ~= 'string' or type(fact.checkerId) ~= 'string' or type(fact.leaseToken) ~= 'string' then
        return nil, 'invalid'
    end
    local raw = redis.call('HGET', tasks_key, fact.taskId)
    if not raw then
        return nil, 'not_found'
    end
    local valid, task = pcall(cjson.decode, raw)
    if not valid or type(task) ~= 'table' then
        return nil, 'invalid'
    end
    if not matches_task(task, fact) then
        return nil, 'observation'
    end
    if task.state == 'DONE' then
        if task.leaseCheckerId == fact.checkerId and task.leaseToken == fact.leaseToken then
            return task, 'ok'
        end
        return nil, 'not_owned'
    end
    if task.state ~= 'LEASED' or tonumber(task.leaseExpiresAtMs or 0) <= now_ms then
        return nil, 'lease_expired'
    end
    if task.leaseCheckerId ~= fact.checkerId or task.leaseToken ~= fact.leaseToken then
        return nil, 'not_owned'
    end
    return task, 'ok'
end

reap()

if operation == 'inflight' then
    return finish({status = 'ok', count = redis.call('ZCARD', queued_key) + redis.call('ZCARD', leases_key)})
end

if operation == 'upstream_inflight' then
    if type(payload.upstreamTasksKey) ~= 'string' or payload.upstreamTasksKey == '' then
        return finish({status = 'invalid'})
    end
    redis.call('ZREMRANGEBYSCORE', payload.upstreamTasksKey, '-inf', now_ms)
    return finish({status = 'ok', count = redis.call('ZCARD', payload.upstreamTasksKey)})
end

if operation == 'due' then
    local result = cjson.decode('[]')
    local scan_limit = tonumber(payload.scanLimit or limit)
    if not scan_limit or scan_limit < limit then
        return finish({status = 'invalid'})
    end
    local ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit)
    for _, proxy_id in ipairs(ids) do
        local _, record = live_record(proxy_id)
        if not record or not valid_state(record.state) then
            redis.call('ZREM', due_key, proxy_id)
		elseif type(payload.upstreamId) == 'string' and payload.upstreamId ~= '' and record.sourceUpstream ~= payload.upstreamId then
            -- The ref remains due for the scheduler that owns this upstream.
        else
            result[#result + 1] = {
                proxyId = proxy_id, upstreamId = record.sourceUpstream, state = record.state, dueAtMs = now_ms
            }
			if #result >= limit then
				break
			end
        end
    end
    local encoded_candidates = cjson.encode(result)
    if #result == 0 then
        encoded_candidates = '[]'
    end
    return finish({status = 'ok', candidatesJSON = encoded_candidates})
end

if operation == 'offer' then
    if type(payload.tasks) ~= 'table' then
        return finish({status = 'invalid'})
    end
    local offered = 0
    for _, task in ipairs(payload.tasks) do
        if type(task) ~= 'table' or task.version ~= 1 or task.level ~= 'BASIC' or task.state ~= 'QUEUED' or
            type(task.taskId) ~= 'string' or type(task.proxyId) ~= 'string' or tonumber(task.deadlineMs or 0) <= now_ms or
            type(task.upstreamId) ~= 'string' or task.upstreamId == '' or type(task.upstreamTasksKey) ~= 'string' or
            task.upstreamTasksKey == '' or tonumber(task.nextDueMs or 0) <= now_ms or tonumber(task.attempts or 0) <= 0 or
            tonumber(task.maxInFlight or 0) <= 0 then
            return finish({status = 'invalid'})
        end
        local score = redis.call('ZSCORE', due_key, task.proxyId)
        local current = redis.call('HGET', ref_task_key, task.proxyId)
        local raw, record = live_record(task.proxyId)
        local active = redis.call('ZCARD', task.upstreamTasksKey)
        if score and tonumber(score) <= now_ms and not current and raw and record.sourceUpstream == task.upstreamId and
            active < tonumber(task.maxInFlight) and valid_state(record.state) then
            if record.state == 'FETCHED' or record.state == 'UNHEALTHY' then
                decrement_state(record.sourceUpstream, record.state)
                increment_state(record.sourceUpstream, 'CHECKING')
                record.state = 'CHECKING'
                raw = cjson.encode(record)
                redis.call('HSET', records_key, task.proxyId, raw)
            end
            redis.call('ZREM', due_key, task.proxyId)
            redis.call('HSET', tasks_key, task.taskId, cjson.encode(task))
            redis.call('HSET', ref_task_key, task.proxyId, task.taskId)
            redis.call('ZADD', queued_key, tonumber(task.priority), task.taskId)
            redis.call('ZADD', task_expiry_key, tonumber(task.deadlineMs), task.taskId)
            redis.call('ZADD', task.upstreamTasksKey, tonumber(task.deadlineMs), task.taskId)
            local expires_at_ms = tonumber(record.expiresAtMs)
            touch(queued_key, expires_at_ms)
            touch(tasks_key, expires_at_ms)
            touch(task_expiry_key, expires_at_ms)
            touch(ref_task_key, expires_at_ms)
			touch(task.upstreamTasksKey, expires_at_ms)
            offered = offered + 1
        end
    end
    return finish({status = 'ok', count = offered})
end

if operation == 'claim' then
    if type(payload.checkerId) ~= 'string' or type(payload.instanceId) ~= 'string' or
        type(payload.levels) ~= 'table' or type(payload.tokens) ~= 'table' then
        return finish({status = 'invalid'})
    end
    local max_in_flight = tonumber(payload.maxInFlight or 0)
    local lease_ttl_ms = tonumber(payload.leaseTTLMS or 0)
    if not max_in_flight or max_in_flight <= 0 or not lease_ttl_ms or lease_ttl_ms <= 0 then
        return finish({status = 'invalid'})
    end
    redis.call('ZREMRANGEBYSCORE', checker_leases_key, '-inf', now_ms)
    local capacity = max_in_flight - redis.call('ZCARD', checker_leases_key)
    if capacity <= 0 then
        return finish({status = 'ok', tasks = cjson.decode('[]')})
    end
    local supported = {}
    for _, level in ipairs(payload.levels) do
        supported[level] = true
    end
    local candidates = redis.call('ZRANGE', queued_key, 0, limit - 1)
    local result = cjson.decode('[]')
    for _, task_id in ipairs(candidates) do
        if #result >= capacity or #result >= #payload.tokens then
            break
        end
        local raw = redis.call('HGET', tasks_key, task_id)
        local valid, task = pcall(cjson.decode, raw or '')
        if not valid or type(task) ~= 'table' or task.state ~= 'QUEUED' or tonumber(task.deadlineMs or 0) <= now_ms then
            remove_task(task_id, true)
        elseif not supported[task.level] then
            -- Preserve unsupported work for a Checker that advertises this level.
        else
            local record_raw, record = live_record(task.proxyId)
            if not record_raw or not record then
                remove_task(task_id, false)
            else
                local lease_expires_at_ms = now_ms + lease_ttl_ms
                if tonumber(task.deadlineMs) < lease_expires_at_ms then
                    lease_expires_at_ms = tonumber(task.deadlineMs)
                end
                local token = payload.tokens[#result + 1]
                if type(token) ~= 'string' or token == '' then
                    return finish({status = 'invalid'})
                end
                task.state = 'LEASED'
                task.leaseCheckerId = payload.checkerId
                task.leaseInstanceId = payload.instanceId
                task.leaseToken = token
                task.leaseExpiresAtMs = lease_expires_at_ms
                task.checkerLeaseKey = checker_leases_key
                local encoded = cjson.encode(task)
                redis.call('HSET', tasks_key, task_id, encoded)
                redis.call('ZREM', queued_key, task_id)
                redis.call('ZADD', leases_key, lease_expires_at_ms, task_id)
                redis.call('ZADD', checker_leases_key, lease_expires_at_ms, task_id)
                touch(checker_leases_key, tonumber(record.expiresAtMs))
                result[#result + 1] = {task = encoded, record = record_raw}
            end
        end
    end
    return finish({status = 'ok', tasks = result})
end

if operation == 'authorize' or operation == 'complete' then
    local task, status = authorize(payload.fact)
    if status ~= 'ok' then
        return finish({status = status})
    end
    if operation == 'complete' and task.state ~= 'DONE' then
        redis.call('ZREM', queued_key, task.taskId)
        redis.call('ZREM', leases_key, task.taskId)
		if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then
			redis.call('ZREM', task.upstreamTasksKey, task.taskId)
		end
        if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
            redis.call('ZREM', task.checkerLeaseKey, task.taskId)
        end
        task.state = 'DONE'
        redis.call('HSET', tasks_key, task.taskId, cjson.encode(task))
        if redis.call('HGET', ref_task_key, task.proxyId) == task.taskId then
            redis.call('HDEL', ref_task_key, task.proxyId)
        end
        local _, record = live_record(task.proxyId)
        if record then
            local next_due_ms = tonumber(task.nextDueMs or 0)
            if next_due_ms <= now_ms then
                next_due_ms = now_ms
            end
            redis.call('ZADD', due_key, next_due_ms, task.proxyId)
            touch(due_key, tonumber(record.expiresAtMs))
        end
    end
    return finish({status = 'ok'})
end

return finish({status = 'invalid'})