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 health_due_key = KEYS[11]
local health_queued_key = KEYS[12]
local health_leases_key = KEYS[13]
local health_tasks_key = KEYS[14]
local health_task_expiry_key = KEYS[15]
local health_ref_task_key = KEYS[16]
local health_unhealthy_key = KEYS[17]

local checked_at_ms = tonumber(ARGV[1])
local next_state = ARGV[2]
local latency_ns = tonumber(ARGV[3])
local cleanup_limit = tonumber(ARGV[4])
local operation_ttl_ms = tonumber(ARGV[5])
local proxy_id = ARGV[6]
local mode = ARGV[7] or 'transition'
local success = ARGV[8]
local max_consecutive_failures = tonumber(ARGV[9])
local task_id = ARGV[10]
local observation_digest = ARGV[11]

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(id, record)
    redis.call('ZREM', available_key, id)
    for _, index_key in ipairs(record and record.indexKeys or {}) do
        redis.call('ZREM', index_key, id)
    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 remove_worker_owned(id)
    local raw = redis.call('HGET', owners_key, 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, id)
    end
end

local function remove_health_task(proxy_id)
    redis.call('ZREM', health_due_key, proxy_id)
    local task_id = redis.call('HGET', health_ref_task_key, proxy_id)
    if not task_id then
        return
    end
    local raw = redis.call('HGET', health_tasks_key, task_id)
    if raw then
        local decoded, task = pcall(cjson.decode, raw)
        if decoded and type(task) == 'table' then
            if type(task.checkerLeaseKey) == 'string' and task.checkerLeaseKey ~= '' then
                redis.call('ZREM', task.checkerLeaseKey, task_id)
            end
            if type(task.upstreamTasksKey) == 'string' and task.upstreamTasksKey ~= '' then
                redis.call('ZREM', task.upstreamTasksKey, task_id)
            end
        end
    end
    redis.call('ZREM', health_queued_key, task_id)
    redis.call('ZREM', health_leases_key, task_id)
    redis.call('ZREM', health_task_expiry_key, task_id)
    redis.call('HDEL', health_tasks_key, task_id)
    redis.call('HDEL', health_ref_task_key, proxy_id)
end

local function sync_worker_owned(id, record)
    local raw = redis.call('HGET', owners_key, 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 > checked_at_ms and
        record.state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms then
        redis.call('ZADD', assignment.workerIndexKey, assignment.expiresAtMs, 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, id)
end

local function remove_proxy(id)
    local raw = redis.call('HGET', records_key, id)
    local record = nil
    if raw then
        record = cjson.decode(raw)
        remove_available(id, record)
        remove_owned(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, 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)
    remove_worker_owned(id)
    redis.call('HDEL', owners_key, id)
    redis.call('ZREM', owner_expiry_key, id)
    redis.call('ZREM', health_unhealthy_key, id)
    remove_health_task(id)
end

local function cleanup_expired()
    local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', checked_at_ms, 'LIMIT', 0, cleanup_limit)
    for _, id in ipairs(expired) do
        remove_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 finish(reply, hard_expiry_ms)
    local encoded = cjson.encode(reply)
    redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
    if hard_expiry_ms then
        local operation_expiry_ms = redis.call('PEXPIRETIME', operation_key)
        if operation_expiry_ms > hard_expiry_ms then
            redis.call('PEXPIREAT', operation_key, hard_expiry_ms)
        end
    end
    return encoded
end

local transitions = {
    FETCHED = {CHECKING = true, EXPIRED = true, REMOVED = true},
    CHECKING = {AVAILABLE = true, UNHEALTHY = true, EXPIRED = true, REMOVED = true},
    AVAILABLE = {SUSPECT = true, DRAINING = true, EXTRACTED = true, EXPIRED = true},
    SUSPECT = {AVAILABLE = true, UNHEALTHY = true, DRAINING = true, EXPIRED = true},
    DRAINING = {EXPIRED = true, UNHEALTHY = true, REMOVED = true},
    UNHEALTHY = {CHECKING = true, REMOVED = true, EXPIRED = true},
    EXTRACTED = {EXPIRED = true, REMOVED = true},
    EXPIRED = {REMOVED = true},
    REMOVED = {},
}

cleanup_expired()

local raw = redis.call('HGET', records_key, proxy_id)
if not raw then
    return finish({status = 'not_found'})
end
local record = cjson.decode(raw)
if type(record.ownerIndexKey) ~= 'string' or record.ownerIndexKey == '' then
    return finish({status = 'invalid'})
end
if tonumber(record.expiresAtMs) <= checked_at_ms then
    remove_proxy(proxy_id)
    return finish({status = 'not_found'})
end

local last_checked_at_ms = tonumber(record.lastCheckedAtMs or '0')
if mode == 'global' then
    if (success ~= '0' and success ~= '1') or not max_consecutive_failures or
        max_consecutive_failures <= 0 or type(task_id) ~= 'string' or task_id == '' or
        type(observation_digest) ~= 'string' or string.len(observation_digest) ~= 64 then
        return finish({status = 'invalid'})
    end
    local last_task_id = record.lastHealthTaskId or ''
    local last_observation_digest = record.lastHealthDigest or ''
    if last_task_id == task_id then
        if last_observation_digest == observation_digest then
            return finish({status = 'ok', record = raw}, tonumber(record.expiresAtMs))
        end
        return finish({status = 'conflict'})
    end
    if checked_at_ms <= last_checked_at_ms then
        return finish({status = 'stale'})
    end
    local consecutive_failures = tonumber(record.consecutiveFailures or '0')
    if not consecutive_failures or consecutive_failures < 0 then
        return finish({status = 'invalid'})
    end
    if success == '1' then
        if record.state ~= 'CHECKING' and record.state ~= 'AVAILABLE' and record.state ~= 'SUSPECT' then
            return finish({status = 'invalid'})
        end
        next_state = 'AVAILABLE'
        consecutive_failures = 0
    else
        consecutive_failures = consecutive_failures + 1
        if record.state == 'CHECKING' then
            next_state = 'UNHEALTHY'
        elseif record.state == 'AVAILABLE' or record.state == 'SUSPECT' then
            if consecutive_failures >= max_consecutive_failures then
                next_state = 'UNHEALTHY'
            else
                next_state = 'SUSPECT'
            end
        else
            return finish({status = 'invalid'})
        end
    end
    record.consecutiveFailures = consecutive_failures
    record.lastHealthTaskId = task_id
    record.lastHealthDigest = observation_digest
    record.lastHealthObservedAtMs = checked_at_ms
else
    if checked_at_ms < last_checked_at_ms then
        return finish({status = 'stale'})
    end
    if checked_at_ms == last_checked_at_ms then
        if record.state ~= next_state then
            return finish({status = 'stale'})
        end
        return finish({status = 'ok', record = raw}, tonumber(record.expiresAtMs))
    end
end
if record.state ~= next_state and not (transitions[record.state] and transitions[record.state][next_state]) then
    return finish({status = 'invalid'})
end

local previous_state = record.state
local was_managed = is_managed(previous_state)
local will_be_managed = is_managed(next_state)
remove_available(proxy_id, record)
record.state = next_state
record.lastCheckedAtMs = checked_at_ms
record.latencyNs = latency_ns
if next_state == 'AVAILABLE' then
    record.lastSuccessAtMs = checked_at_ms
end
if next_state == 'UNHEALTHY' then
    local unhealthy_since_ms = tonumber(record.unhealthySinceMs or '0')
    if not unhealthy_since_ms or unhealthy_since_ms <= 0 then
        unhealthy_since_ms = checked_at_ms
    end
    record.unhealthySinceMs = unhealthy_since_ms
    redis.call('ZADD', health_unhealthy_key, unhealthy_since_ms, proxy_id)
else
    record.unhealthySinceMs = nil
    redis.call('ZREM', health_unhealthy_key, proxy_id)
end
if was_managed and not will_be_managed then
    decrement_inventory(record.sourceUpstream)
elseif not was_managed and will_be_managed then
    redis.call('HINCRBY', inventory_key, record.sourceUpstream, 1)
end
if previous_state ~= next_state then
    decrement_state(record.sourceUpstream, previous_state)
    increment_state(record.sourceUpstream, next_state)
end

local encoded = cjson.encode(record)
redis.call('HSET', records_key, proxy_id, encoded)
local owned = (record.ownerWorkerId and record.ownerWorkerId ~= '') or redis.call('HEXISTS', owners_key, proxy_id) == 1
remove_owned(proxy_id, record)
if next_state == 'AVAILABLE' and tonumber(record.usableUntilMs) > checked_at_ms then
    if owned then
        redis.call('ZADD', record.ownerIndexKey, record.usableUntilMs, proxy_id)
        touch(record.ownerIndexKey, tonumber(record.expiresAtMs))
    else
        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
end
sync_worker_owned(proxy_id, record)
touch(records_key, tonumber(record.expiresAtMs))
touch(unique_key, tonumber(record.expiresAtMs))
touch(idkeys_key, tonumber(record.expiresAtMs))
touch(expiry_key, tonumber(record.expiresAtMs))
touch(available_key, tonumber(record.expiresAtMs))
touch(inventory_key, tonumber(record.expiresAtMs))
touch(state_inventory_key, tonumber(record.expiresAtMs))
touch(owners_key, tonumber(record.expiresAtMs))
touch(owner_expiry_key, tonumber(record.expiresAtMs))

return finish({status = 'ok', record = encoded}, tonumber(record.expiresAtMs))