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 egress_due_key = KEYS[11]
local target_due_key = KEYS[12]

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 due_key_for_level(level)
    if level == 'EGRESS' then
        return egress_due_key
    end
    if level == 'TARGET' then
        return target_due_key
    end
    return due_key
end

local function task_reference(level, proxy_id, routing_name, target_url)
    if level == 'EGRESS' then
        return proxy_id .. '\0EGRESS\0\0' .. target_url
    end
    if level == 'TARGET' then
        return proxy_id .. '\0TARGET\0' .. routing_name .. '\0' .. target_url
    end
    return proxy_id
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
        local reference = task.reference or task.proxyId
        if type(reference) == 'string' and redis.call('HGET', ref_task_key, reference) == task_id then
            redis.call('HDEL', ref_task_key, reference)
        end
		local due_member = task.dueMember or task.proxyId
		if requeue and type(task.proxyId) == 'string' and type(due_member) == 'string' then
            local _, record = live_record(task.proxyId)
            if record then
				local task_due_key = due_key_for_level(task.level)
                redis.call('ZADD', task_due_key, now_ms, due_member)
                touch(task_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)
	if not task or task.taskId ~= fact.taskId or task.proxyId ~= fact.proxyId or task.level ~= fact.level then
		return false
	end
	if task.level == 'EGRESS' then
		return (fact.routingName or '') == '' and (fact.targetUrl or '') == ''
	end
    return (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
    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 level = payload.level or 'BASIC'
	if level ~= 'BASIC' and level ~= 'EGRESS' and level ~= 'TARGET' then
		return finish({status = 'invalid'})
	end
	if level == 'EGRESS' and (type(payload.targetUrl) ~= 'string' or payload.targetUrl == '') then
		return finish({status = 'invalid'})
	end
	if level == 'TARGET' and (type(payload.routingName) ~= 'string' or payload.routingName == '' or
		type(payload.targetUrl) ~= 'string' or payload.targetUrl == '') then
		return finish({status = 'invalid'})
	end
	local result_members = {}
	local function append_candidate(proxy_id, record, member)
		if result_members[member] then
			return
		end
		result_members[member] = true
		result[#result + 1] = {
			proxyId = proxy_id, upstreamId = record.sourceUpstream, state = record.state, level = level,
			routingName = level == 'TARGET' and payload.routingName or '',
			targetUrl = (level == 'EGRESS' or level == 'TARGET') and payload.targetUrl or '', dueAtMs = now_ms
		}
	end
	if level == 'EGRESS' or level == 'TARGET' then
		local basic_ids = redis.call('ZRANGEBYSCORE', due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit)
		for _, proxy_id in ipairs(basic_ids) do
			local _, record = live_record(proxy_id)
			if record and valid_state(record.state) and
				(type(payload.upstreamId) ~= 'string' or payload.upstreamId == '' or record.sourceUpstream == payload.upstreamId) then
				local member = payload.targetUrl .. '\0' .. proxy_id
				if level == 'TARGET' then
					member = payload.routingName .. '\0' .. payload.targetUrl .. '\0' .. proxy_id
				end
				local reference = task_reference(level, proxy_id, payload.routingName or '', payload.targetUrl)
				if not redis.call('HGET', ref_task_key, reference) then
					local task_due_key = due_key_for_level(level)
					redis.call('ZADD', task_due_key, now_ms, member)
					touch(task_due_key, tonumber(record.expiresAtMs))
				end
			end
		end
	end
	local source_due_key = due_key_for_level(level)
	local ids = redis.call('ZRANGEBYSCORE', source_due_key, '-inf', now_ms, 'LIMIT', 0, scan_limit)
	for _, member in ipairs(ids) do
		local proxy_id = member
		local matching_target = true
		if level == 'EGRESS' then
			local separator = string.find(member, '\0', 1, true)
			if not separator or string.sub(member, 1, separator - 1) ~= payload.targetUrl then
				matching_target = false
			else
				proxy_id = string.sub(member, separator + 1)
			end
		elseif level == 'TARGET' then
			local first_separator = string.find(member, '\0', 1, true)
			local second_separator = first_separator and string.find(member, '\0', first_separator + 1, true) or nil
			if not first_separator or not second_separator or string.sub(member, 1, first_separator - 1) ~= payload.routingName or
				string.sub(member, first_separator + 1, second_separator - 1) ~= payload.targetUrl then
				matching_target = false
			else
				proxy_id = string.sub(member, second_separator + 1)
			end
		end
		if matching_target then
			local _, record = live_record(proxy_id)
			local reference = task_reference(level, proxy_id, payload.routingName or '', payload.targetUrl or '')
			if not record or not valid_state(record.state) then
				redis.call('ZREM', source_due_key, member)
			elseif (level == 'EGRESS' or level == 'TARGET') and redis.call('HGET', ref_task_key, reference) then
				redis.call('ZREM', source_due_key, member)
			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
				append_candidate(proxy_id, record, member)
				if #result >= limit then
					break
				end
			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' and task.level ~= 'EGRESS' and task.level ~= 'TARGET') 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 type(task.reference) ~= 'string' or task.reference == '' or
			type(task.dueMember) ~= 'string' or task.dueMember == '' 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
		if task.level == 'BASIC' and ((task.routingName or '') ~= '' or (task.targetUrl or '') ~= '') then
			return finish({status = 'invalid'})
		end
		if task.level == 'EGRESS' and ((task.routingName or '') ~= '' or type(task.targetUrl) ~= 'string' or task.targetUrl == '') then
			return finish({status = 'invalid'})
		end
		if task.level == 'TARGET' and (type(task.routingName) ~= 'string' or task.routingName == '' or
			type(task.targetUrl) ~= 'string' or task.targetUrl == '') then
			return finish({status = 'invalid'})
		end
		local task_due_key = due_key_for_level(task.level)
        local score = redis.call('ZSCORE', task_due_key, task.dueMember)
        local current = redis.call('HGET', ref_task_key, task.reference)
        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 task.level == 'BASIC' and (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', task_due_key, task.dueMember)
            redis.call('HSET', tasks_key, task.taskId, cjson.encode(task))
			redis.call('HSET', ref_task_key, task.reference, 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))
		local reference = task.reference or task.proxyId
        if redis.call('HGET', ref_task_key, reference) == task.taskId then
            redis.call('HDEL', ref_task_key, reference)
        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
			local task_due_key = due_key_for_level(task.level)
            redis.call('ZADD', task_due_key, next_due_ms, task.dueMember or task.proxyId)
            touch(task_due_key, tonumber(record.expiresAtMs))
        end
    end
    return finish({status = 'ok'})
end

return finish({status = 'invalid'})