z_mirror/Krastorio2

git clone https://git.raiguard.me/z_mirror/Krastorio2

Caleb HeuerChange to four-space indentation07409ca62026-01-05

local flib_bounding_box = require("__flib__.bounding-box")

--- @alias k2.ResourceRestrictions.Forces table<ForceIndex, k2.ResourceRestrictions.ForceData?>

--- @class k2.ResourceRestrictions.ForceData
--- @field disabled_mining_drills table<UnitNumber, k2.ResourceRestrictions.DrillData>
--- @field force LuaForce
--- @field restricted_resources Set<string>

--- @class k2.ResourceRestrictions.DrillData
--- @field drill LuaEntity
--- @field icon LuaRenderObject

--- @return k2.ResourceRestrictions.Forces
local function get_forces()
    if not storage.resource_restrictions then
        --- @type k2.ResourceRestrictions.Forces
        storage.resource_restrictions = {}
    end
    return storage.resource_restrictions
end

--- @param force LuaForce
local function get_force_data(force)
    local forces = get_forces()
    local force_index = force.index
    local force_data = forces[force_index]
    if not force_data then
        force_data = {
            disabled_mining_drills = {},
            force = force,
            restricted_resources = {},
        }
        forces[force_index] = force_data
    end
    return force_data
end

--- @param force_data k2.ResourceRestrictions.ForceData
--- @param drill LuaEntity
--- @return LuaEntityPrototype? restricted_resource
local function get_first_restricted_resource(force_data, drill)
    local radius = drill.prototype.get_mining_drill_radius(drill.quality)
    --- @cast radius -?
    local area = flib_bounding_box.from_dimensions(drill.position, radius * 2, radius * 2)
    local targets = drill.surface.find_entities_filtered({ type = "resource", area = area })
    local drill_categories = drill.prototype.resource_categories or {}
    local restrictions = force_data.restricted_resources
    for i = 1, #targets do
        local target = targets[i]
        if drill_categories[target.prototype.resource_category] and restrictions[target.name] then
            return target.prototype
        end
    end
end

--- @param force_data k2.ResourceRestrictions.ForceData
--- @param drill LuaEntity
local function update_drill(force_data, drill)
    local disabled_mining_drills = force_data.disabled_mining_drills

    local unit_number = drill.unit_number
    --- @cast unit_number -?

    local drill_data = disabled_mining_drills[unit_number]

    local restriction = get_first_restricted_resource(force_data, drill)
    if restriction then
        if not drill_data then
            drill.active = false
            drill.custom_status = {
                diode = defines.entity_status_diode.red,
                label = { "entity-status.kr-unable-to-mine", restriction.localised_name },
            }
            disabled_mining_drills[unit_number] = {
                drill = drill,
                icon = rendering.draw_sprite({
                    sprite = "utility/resources_depleted_icon",
                    target = drill,
                    surface = drill.surface,
                    x_scale = 0.5,
                    y_scale = 0.5,
                    blink_interval = 30,
                }),
            }
        end
    else
        if drill_data then
            drill.active = true
            drill.custom_status = nil
            if drill_data.icon.valid then
                drill_data.icon.destroy()
            end
            disabled_mining_drills[unit_number] = nil
        end
    end
end

--- @param force_data k2.ResourceRestrictions.ForceData
local function refresh_all_drills(force_data)
    for _, drill_data in pairs(force_data.disabled_mining_drills) do
        if drill_data.icon.valid then
            drill_data.icon.destroy()
        end
    end
    force_data.disabled_mining_drills = {}
    for _, surface in pairs(game.surfaces) do
        for _, drill in pairs(surface.find_entities_filtered({ type = "mining-drill", force = force_data.force })) do
            update_drill(force_data, drill)
        end
    end
end

local unlocks = prototypes.mod_data["kr-resource-restrictions"].data --[[ @as table<data.TechnologyID, data.EntityID[]?> ]]

--- @param force_data k2.ResourceRestrictions.ForceData
--- @param technology LuaTechnology
--- @param restrict boolean
local function process_effects(force_data, technology, restrict)
    local technology_unlocks = unlocks[technology.name]
    if not technology_unlocks then
        return
    end

    local restricted_resources = force_data.restricted_resources
    for _, resource_name in pairs(technology_unlocks) do
        restricted_resources[resource_name] = restrict or nil
    end

    if restrict then
        -- The on_research_reversed event is very rare, and never happens during regular gameplay, so it is acceptable to have a lag spike in this case.
        refresh_all_drills(force_data)
    else
        for _, drill_data in pairs(force_data.disabled_mining_drills) do
            update_drill(force_data, drill_data.drill)
        end
    end
end

--- @param force LuaForce
local function refresh_force(force)
    force.mining_with_fluid = true
    local force_data = get_force_data(force)
    for _, technology in pairs(force.technologies) do
        if not technology.researched then
            -- TODO: This is doing a full scan of all surfaces every time!
            process_effects(force_data, technology, true)
        end
    end
end

--- @param e EntityBuiltEvent
local function on_entity_built(e)
    local entity = e.entity or e.destination
    if not entity.valid or entity.type ~= "mining-drill" then
        return
    end
    update_drill(get_force_data(entity.force --[[@as LuaForce]]), entity)
end

--- @param e EntityDestroyedEvent
local function on_entity_destroyed(e)
    local entity = e.entity
    if not entity.valid or entity.type ~= "mining-drill" then
        return
    end

    -- The icon will get destroyed automatically because its target is the mining drill.
    get_force_data(e.entity.force --[[@as LuaForce]]).disabled_mining_drills[entity.unit_number] = nil
end

--- @param e EventData.on_research_finished
local function on_research_finished(e)
    process_effects(get_force_data(e.research.force), e.research, false)
end

--- @param e EventData.on_research_finished
local function on_research_reversed(e)
    process_effects(get_force_data(e.research.force), e.research, true)
end

--- @param e EventData.on_force_created
local function on_force_created(e)
    refresh_force(e.force)
end

--- @param e EventData.on_force_reset
local function on_force_reset(e)
    refresh_force(e.force)
end

--- @param e EventData.on_forces_merging
local function on_forces_merging(e)
    local source_data = get_force_data(e.source)
    for _, drill_data in pairs(source_data.disabled_mining_drills) do
        if drill_data.icon.valid then
            drill_data.icon.destroy()
        end
    end
    get_forces()[e.source.index] = nil
    refresh_force(e.destination)
end

--- @param e EventData.on_technology_effects_reset
local function on_technology_effects_reset(e)
    refresh_force(e.force)
end

local function refresh()
    storage.resource_restrictions = nil
    for _, force in pairs(game.forces) do
        refresh_force(force)
        refresh_all_drills(get_force_data(force))
    end
end

--- @class MiningDrillRestrictionsHandler: event_handler
local handler = {}

handler.on_init = refresh
handler.on_configuration_changed = refresh

handler.events = {
    [defines.events.on_built_entity] = on_entity_built,
    [defines.events.on_entity_cloned] = on_entity_built,
    [defines.events.on_entity_died] = on_entity_destroyed,
    [defines.events.on_force_created] = on_force_created,
    [defines.events.on_force_reset] = on_force_reset,
    [defines.events.on_forces_merging] = on_forces_merging,
    [defines.events.on_player_mined_entity] = on_entity_destroyed,
    [defines.events.on_research_finished] = on_research_finished,
    [defines.events.on_research_reversed] = on_research_reversed,
    [defines.events.on_robot_built_entity] = on_entity_built,
    [defines.events.on_robot_mined_entity] = on_entity_destroyed,
    [defines.events.on_space_platform_built_entity] = on_entity_built,
    [defines.events.on_space_platform_mined_entity] = on_entity_destroyed,
    [defines.events.on_technology_effects_reset] = on_technology_effects_reset,
    [defines.events.script_raised_built] = on_entity_built,
    [defines.events.script_raised_destroy] = on_entity_destroyed,
    [defines.events.script_raised_revive] = on_entity_built,
}

return handler