lua-prometheus-exporter/power-stats.lua

function AddPowerPole(event)
    local e = event.entity
    if e then
        storage.representativePoles[e.unit_number] = e
        -- Update cache with new network
        storage.networkCache = nil
    end
end

function RemovePowerPole(event)
    local e = event.entity
    if e then
        storage.representativePoles[e.unit_number] = nil
        -- Invalidate cache - force rebuild to find new representative if needed
        storage.networkCache = nil
    end
end

function GetNetworks()
    -- Return cached networks if available, otherwise rebuild
    if storage.networkCache then
        return storage.networkCache
    end

    local networks = {}
    for _, pole in pairs(storage.representativePoles) do
        if pole.valid then
            local netID = pole.electric_network_id
            if netID then
                networks[netID] = pole
            end
        end
    end

    storage.networkCache = networks
    return networks
end

function ScanNetworks()
    storage.representativePoles = {}
    storage.networkCache = {}

    for _, surface in pairs(game.surfaces) do
        for _, pole in pairs(surface.find_entities_filtered { type = "electric-pole" }) do
            if pole.valid and pole.electric_network_id then
                storage.representativePoles[pole.unit_number] = pole
                storage.networkCache[pole.electric_network_id] = pole
            end
        end
    end
end

solarBase = 60
steamBase = 900
turbineBase = 5820
fusionBase = 50000


function AddGenerator(event)
    if event then
        if event.entity then
            ---@type LuaEntity
            local entity = event.entity
            if entity.type == "solar-panel" then
                storage.powerStats["solar"][entity.surface.name] = (storage.powerStats["solar"][entity.surface.name] or 0) +
                (solarBase * entity.quality.default_multiplier * entity.surface.solar_power_multiplier)
            end
            if entity.type == "fusion-generator" then
                storage.powerStats["fusion"][entity.surface.name] = (storage.powerStats["fusion"][entity.surface.name] or 0) +
                (fusionBase * entity.quality.default_multiplier)
            end
            if entity.name == "steam-engine" then
                storage.powerStats["steam"][entity.surface.name] = (storage.powerStats["steam"][entity.surface.name] or 0) +
                (steamBase * entity.quality.default_multiplier)
            end
            if entity.name == "steam-turbine" then
                storage.powerStats["turbine"][entity.surface.name] = (storage.powerStats["turbine"][entity.surface.name] or 0) +
                (turbineBase * entity.quality.default_multiplier)
            end
        end
    end
end

function RemoveGenerator(event)
    if event then
        if event.entity then
            ---@type LuaEntity
            local entity = event.entity
            if entity.type == "solar-panel" then
                storage.powerStats["solar"][entity.surface.name] = (storage.powerStats["solar"][entity.surface.name] or 0) -
                (solarBase * entity.quality.default_multiplier * entity.surface.solar_power_multiplier)
            end
            if entity.type == "fusion-generator" then
                storage.powerStats["fusion"][entity.surface.name] = (storage.powerStats["fusion"][entity.surface.name] or 0) -
                (fusionBase * entity.quality.default_multiplier)
            end
            if entity.name == "steam-engine" then
                storage.powerStats["steam"][entity.surface.name] = (storage.powerStats["steam"][entity.surface.name] or 0) -
                (steamBase * entity.quality.default_multiplier)
            end
            if entity.name == "steam-turbine" then
                storage.powerStats["turbine"][entity.surface.name] = (storage.powerStats["turbine"][entity.surface.name] or 0) -
                (turbineBase * entity.quality.default_multiplier)
            end
        end
    end
end

function ScanGenerators()
    storage.powerGenerators = {}
    for _, surface in pairs(game.surfaces) do
        for _, generator in pairs(surface.find_entities_filtered({ filter = type, type = { "generator", "fusion-generator", "solar-panel" } })) do
            storage.powerGenerators[generator.unit_number] = generator
        end
    end
    PopulateStats()
end

function PopulateStats()
    local solar = {}
    local steam = {}
    local turbine = {}
    local fusion = {}

    for _, generator in pairs(storage.powerGenerators) do
        local surface = generator.surface
        if generator.type == "solar-panel" then
            local surfaceSolarFactor = surface.solar_power_multiplier
            solar[surface.name] = (solar[surface.name] or 0) +
            (solarBase * generator.quality.default_multiplier * surfaceSolarFactor)
        end
        if generator.type == "fusion-generator" then
            fusion[surface.name] = (fusion[surface.name] or 0) + (fusionBase * generator.quality.default_multiplier)
        end
        if generator.name == "steam-engine" then
            steam[surface.name] = (steam[surface.name] or 0) + (steamBase * generator.quality.default_multiplier)
        end
        if generator.name == "steam-turbine" then
            turbine[surface.name] = (turbine[surface.name] or 0) + (turbineBase * generator.quality.default_multiplier)
        end
    end
    storage.powerStats["solar"] = solar
    storage.powerStats["steam"] = steam
    storage.powerStats["turbine"] = turbine
    storage.powerStats["fusion"] = fusion
end

function GetPossiblePower()
    local result = {}
    result[#result + 1] = "---max-power---\n"
    for _, surface in pairs(game.surfaces) do
        result[#result + 1] = ("%s:%d:%d:%d:%d"):format(
            surface.name,
            storage.powerStats["solar"][surface.name] or 0,
            storage.powerStats["steam"][surface.name] or 0,
            storage.powerStats["turbine"][surface.name] or 0,
            storage.powerStats["fusion"][surface.name] or 0
        )
    end
    return table.concat(result, "\n")
end


function SendPowerStats()
    if options.enablePower then
        local possiblePower = GetPossiblePower()
        local powerPart = {}
        powerPart[#powerPart + 1] = "---power-stats---\n"
        for _, pole in pairs(GetNetworks()) do
            if pole.valid and pole.type == "electric-pole" then
                local input = pole.electric_network_statistics.input_counts
                local output = pole.electric_network_statistics.output_counts
                local capacity = pole.electric_network_statistics.storage_counts
                local surfaceName = pole.surface.name
                for item, value in pairs(input) do
                    powerPart[#powerPart + 1] = ("%s:%d:in:%s:%d"):format(surfaceName, pole.electric_network_id, item,
                        value)
                end
                for item, value in pairs(output) do
                    powerPart[#powerPart + 1] = ("%s:%d:out:%s:%d"):format(surfaceName, pole.electric_network_id, item,
                        value)
                end
                for item, value in pairs(capacity) do
                    powerPart[#powerPart + 1] = ("%s:%d:capacity:%s:%d"):format(surfaceName, pole.electric_network_id,
                        item, value)
                end
            end
        end
        powerPart[#powerPart + 1] = possiblePower
        helpers.send_udp(udpAddress, table.concat(powerPart, "\n"), serverIndex)
    end
end