z_mirror/Krastorio2

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

Caleb HeuerChange to four-space indentation07409ca62026-01-05

--- @class BeamData
--- @field beam LuaEntity
--- @field beam_number uint64
--- @field target_data TargetData
--- @field tower_data TowerData

--- @class GridData
--- @field absorber LuaEquipment?
--- @field grid LuaEquipmentGrid

--- @class TargetData
--- @field connections TargetConnections
--- @field entity LuaEntity
--- @field full_tick uint64?
--- @field grid_data GridData
--- @field unit_number uint64

--- @class TargetConnections
--- @field by_beam table<number, ConnectionData>
--- @field by_tower table<number, ConnectionData>

--- @class TowerData
--- @field entities TowerEntities
--- @field tower_unit_number uint64
--- @field turret_unit_number uint64

--- @class TowerEntities
--- @field collision LuaEntity
--- @field tower LuaEntity
--- @field turret LuaEntity

local charging_rate = 3000000 -- 3 MW
local cooldown = 10
local loss_multiplier = 1.8
local range = 20
local required_energy = 10000000 -- 10 MW

local absorber_buffer_capacity = prototypes.equipment["kr-energy-absorber-equipment"].energy_source.buffer_capacity

-- TOWER
-- The entity that is interacted with

--- @param e EntityBuiltEvent
local function on_entity_built(e)
    local entity = e.entity or e.destination
    if not entity or not entity.valid then
        return
    end
    local entity_name = entity.name

    -- Internal entities should not be cloned
    if
        e.name == defines.events.on_entity_cloned
        and (entity_name == "kr-tesla-coil-turret" or entity_name == "kr-tesla-coil-collision")
    then
        entity.destroy()
        return
    end

    if entity_name ~= "kr-tesla-coil" then
        return
    end

    local surface = entity.surface
    local unit_number = entity.unit_number --[[@as uint]]

    local turret = surface.create_entity({
        name = "kr-tesla-coil-turret",
        position = entity.position,
        force = game.forces["kr-internal-turrets"],
        create_build_effect_smoke = false,
        raise_built = true,
    })
    if not turret or not turret.valid then
        game.print(
            "Building tesla failed due to AAI Programmable Vehicles; This tesla coil will not function. To fix this, change the 'Deadzone construction denial range' map setting to zero."
        )
        entity.active = false
        return
    end
    turret.destructible = false

    local collision_entity = surface.create_entity({
        name = "kr-tesla-coil-collision",
        position = entity.position,
        force = entity.force,
        create_build_effect_smoke = false,
        raise_built = true,
    })
    assert(collision_entity, "Tesla coil collision entity was destroyed")
    collision_entity.destructible = false

    --- @type TowerData
    local data = {
        entities = {
            collision = collision_entity,
            tower = entity,
            turret = turret,
        },
        tower_unit_number = unit_number,
        turret_unit_number = turret.unit_number,
    }

    storage.tesla_coil.turrets[data.turret_unit_number] = data
    storage.tesla_coil.towers[unit_number] = data
end

--- @param entity LuaEntity
local function destroy_if_valid(entity)
    if entity.valid then
        entity.destroy()
    end
end

--- @param e EntityDestroyedEvent
local function on_entity_destroyed(e)
    local entity = e.entity
    if not entity or not entity.valid or entity.name ~= "kr-tesla-coil" then
        return
    end

    local unit_number = entity.unit_number
    assert(unit_number)
    local tower_data = storage.tesla_coil.towers[unit_number]
    if not tower_data then
        return
    end

    destroy_if_valid(tower_data.entities.collision)
    destroy_if_valid(tower_data.entities.turret)

    storage.tesla_coil.turrets[tower_data.turret_unit_number] = nil
    storage.tesla_coil.towers[unit_number] = nil
end

--- @param grid LuaEquipmentGrid
--- @param equipment_name string
--- @return LuaEquipment?
local function grid_find_any(grid, equipment_name)
    for _, equipment in pairs(grid.equipment) do
        if equipment.name == equipment_name then
            return equipment
        end
    end
end

-- TARGET
-- An entity that will receive energy from a tesla coil

--- @param target LuaEntity
--- @return GridData?
local function get_grid_data(target)
    --- @type LuaEquipmentGrid
    local grid
    if target.type == "character" then
        local armor_inventory = target.get_inventory(defines.inventory.character_armor)
        if armor_inventory and armor_inventory.valid then
            local armor = armor_inventory[1]
            if armor and armor.valid_for_read then
                grid = armor.grid
            end
        end
    else
        grid = target.grid
    end

    if grid then
        return {
            absorber = grid_find_any(grid, "kr-energy-absorber-equipment"),
            grid = grid,
        }
    end
end

--- Updates the absorber object in a target's equipment grid
--- @param grid LuaEquipmentGrid
local function update_target_grid(grid)
    for _, target_data in pairs(storage.tesla_coil.targets) do
        local grid_data = target_data.grid_data
        if grid_data.grid.valid and grid_data.grid == grid then
            grid_data.absorber = grid_find_any(grid, "kr-energy-absorber-equipment")
        end
    end
end

--- @param target LuaEntity
--- @param tower_data TowerData
local function add_target(target, tower_data)
    local target_unit_number = target.unit_number
    --- @cast target_unit_number -?
    local tower = tower_data.entities.tower
    if not tower.valid or tower.energy < required_energy then
        return
    end
    -- Check the target's equipment grid for an energy absorber
    local grid_data = get_grid_data(target)
    if grid_data and grid_data.absorber and grid_data.absorber.valid then
        --- @type TargetData
        local target_data = {
            connections = {
                --- @type table<number, ConnectionData>
                by_beam = {},
                --- @type table<number, ConnectionData>
                by_tower = {},
            },
            entity = target,
            --- @type number?
            full_tick = nil,
            grid_data = grid_data,
            unit_number = target_unit_number,
        }
        storage.tesla_coil.targets[target_unit_number] = target_data
        return target_data
    end
end

--- @param target_unit_number number
local function remove_target(target_unit_number)
    storage.tesla_coil.targets[target_unit_number] = nil
end

-- CONNECTION
-- A connection between a tower and a target, comprising of a beam
-- There can be unlimited connections per target

--- @param target_data TargetData
--- @param tower_data TowerData
--- @return boolean
local function add_connection(target_data, tower_data)
    -- Check if the absorber has space
    local absorber = target_data.grid_data.absorber
    if not absorber or not absorber.valid then
        remove_target(target_data.unit_number)
        return false
    end

    local tower = tower_data.entities.tower
    if not tower.valid then
        remove_target(target_data.unit_number)
        return false
    end

    if absorber.energy >= absorber_buffer_capacity then
        return false
    end

    -- Create beam entity
    local beam = tower_data.entities.tower.surface.create_entity({
        name = "kr-tesla-coil-electric-beam",
        source = tower_data.entities.tower,
        source_offset = { 0, -2.2 },
        position = tower_data.entities.tower.position,
        target = target_data.entity,
        duration = 0,
        max_length = range,
        force = tower_data.entities.tower.force,
        raise_built = true,
    })
    if not beam then
        return false
    end
    local beam_number = script.register_on_object_destroyed(beam)

    storage.tesla_coil.beams[beam_number] = {
        beam = beam,
        beam_number = beam_number,
        target_data = target_data,
        tower_data = tower_data,
    }

    --- @class ConnectionData
    local connection_data = {
        beam = beam,
        beam_number = beam_number,
        tower_data = tower_data,
    }
    target_data.connections.by_beam[beam_number] = connection_data
    target_data.connections.by_tower[tower_data.tower_unit_number] = connection_data

    return true
end

--- @param target_data TargetData
--- @param tower_data TowerData
local function remove_connection(target_data, tower_data)
    local connection_data = target_data.connections.by_tower[tower_data.tower_unit_number]

    -- Destroy beam if it still exists
    if connection_data.beam.valid then
        connection_data.beam.destroy()
    end

    local beam_number = connection_data.beam_number
    storage.tesla_coil.beams[beam_number] = nil

    target_data.connections.by_beam[beam_number] = nil
    target_data.connections.by_tower[tower_data.tower_unit_number] = nil

    if table_size(target_data.connections.by_beam) == 0 then
        remove_target(target_data.unit_number)
    end
end

--- @param target_data TargetData
--- @param tower_data TowerData
local function update_connection(target_data, tower_data)
    local absorber = target_data.grid_data.absorber

    -- Check if the tower is powered
    if not absorber or not absorber.valid or tower_data.entities.tower.energy < required_energy then
        remove_connection(target_data, tower_data)
        return
    end

    local energy = absorber.energy
    if energy < absorber_buffer_capacity then
        -- Calculate how much to add
        local to_add = charging_rate / 60 * cooldown
        local result = energy + to_add
        local tower = tower_data.entities.tower

        if result >= absorber_buffer_capacity then
            absorber.energy = absorber_buffer_capacity
            target_data.full_tick = game.tick
        else
            absorber.energy = result
            target_data.full_tick = nil
        end

        tower.energy = tower.energy - (to_add * loss_multiplier)
    elseif target_data.full_tick and target_data.full_tick + cooldown <= game.tick then
        remove_connection(target_data, tower_data)
    end
end

--- @param target LuaEntity
--- @param turret LuaEntity
local function process_turret_fire(target, turret)
    local tower_data = storage.tesla_coil.turrets[turret.unit_number]
    if not tower_data then
        return
    end

    local target_data = storage.tesla_coil.targets[target.unit_number]
    if not target_data then
        target_data = add_target(target, tower_data)
    end

    if target_data then
        -- Just in case (#182)
        if not target_data.entity.valid then
            target_data.entity = target
        end
        local connection = target_data.connections.by_tower[tower_data.tower_unit_number]
        if connection or add_connection(target_data, tower_data) then
            update_connection(target_data, tower_data)
        end
    end
end

-- EVENTS

--- @param e EventData.on_object_destroyed
local function on_object_destroyed(e)
    local beam_data = storage.tesla_coil.beams[e.registration_number]
    if beam_data then
        remove_connection(beam_data.target_data, beam_data.tower_data)
    end
end

--- @param e EventData.on_equipment_inserted
local function on_equipment_inserted(e)
    if not e.equipment.valid or e.equipment.name ~= "kr-energy-absorber-equipment" then
        return
    end
    update_target_grid(e.grid)
end

--- @param e EventData.on_equipment_removed
local function on_equipment_removed(e)
    if e.equipment == "kr-energy-absorber-equipment" then
        update_target_grid(e.grid)
    end
end

--- @param e EventData.on_script_trigger_effect
local function on_script_trigger_effect(e)
    if e.effect_id == "kr-tesla-coil-trigger" then
        process_turret_fire(e.target_entity, e.source_entity)
    end
end

--- @class TeslaCoilHandler : event_handler
local tesla_coil = {}

function tesla_coil.on_init()
    storage.tesla_coil = {
        --- @type table<number, BeamData>
        beams = {},
        --- @type table<number, TowerData>
        turrets = {},
        --- @type table<number, TowerData>
        towers = {},
        --- @type table<number, TargetData>
        targets = {},
    }
end

tesla_coil.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_equipment_inserted] = on_equipment_inserted,
    [defines.events.on_equipment_removed] = on_equipment_removed,
    [defines.events.on_object_destroyed] = on_object_destroyed,
    [defines.events.on_player_mined_entity] = on_entity_destroyed,
    [defines.events.on_robot_built_entity] = on_entity_built,
    [defines.events.on_robot_mined_entity] = on_entity_destroyed,
    [defines.events.on_script_trigger_effect] = on_script_trigger_effect,
    [defines.events.script_raised_built] = on_entity_built,
    [defines.events.script_raised_destroy] = on_entity_destroyed,
    [defines.events.script_raised_revive] = on_entity_built,
    [defines.events.on_space_platform_built_entity] = on_entity_built,
    [defines.events.on_space_platform_mined_entity] = on_entity_destroyed,
}

return tesla_coil

--[[
  -- Quick testing code
  event.register(remote.call("EditorExtensions", "debug_world_ready_event"), function()
    local surface = game.surfaces.nauvis

    local coils = {
      { x = -9.5, y = -9.5 },
      { x = 10.5, y = -9.5 },
      { x = -9.5, y = 10.5 },
      { x = 10.5, y = 10.5 },
    }
    for _, position in pairs(coils) do
      surface.create_entity({
        name = "kr-tesla-coil",
        direction = defines.direction.north,
        position = position,
        force = game.forces.player,
        raise_built = true,
        create_build_effect_smoke = false,
      })
    end
    surface.create_entity({
      name = "ee-infinity-accumulator-primary-output",
      position = { x = -1, y = -9 },
      force = game.forces.player,
      raise_built = true,
      create_build_effect_smoke = false,
    })
    surface.create_entity({
      name = "ee-super-substation",
      position = { x = 1, y = -9 },
      force = game.forces.player,
      raise_built = true,
      create_build_effect_smoke = false,
    })

    local character = game.get_player(1).character
    local armor_inventory = character.get_inventory(defines.inventory.character_armor)
    if armor_inventory and armor_inventory.valid then
      local armor = armor_inventory[1]
      if armor and armor.valid_for_read then
        armor.grid.put({ name = "kr-energy-absorber-equipment", position = { 5, 5 } })
      end
    end
  end)
]]