tbd-station-14/Content.Server/Power/EntitySystems/PowerNetSystem.cs

#nullable enable
using System.Collections.Generic;
using Content.Server.NodeContainer.EntitySystems;
using Content.Server.Power.Components;
using Content.Server.Power.NodeGroups;
using Content.Server.Power.Pow3r;
using JetBrains.Annotations;
using Robust.Shared.GameObjects;
using Robust.Shared.Maths;

namespace Content.Server.Power.EntitySystems
{
    /// <summary>
    ///     Manages power networks, power state, and all power components.
    /// </summary>
    [UsedImplicitly]
    public class PowerNetSystem : EntitySystem
    {
        private readonly PowerState _powerState = new();
        private readonly HashSet<PowerNet> _powerNetReconnectQueue = new();
        private readonly HashSet<ApcNet> _apcNetReconnectQueue = new();

        private int _nextId = 1;
        private readonly BatteryRampPegSolver _solver = new();

        public override void Initialize()
        {
            base.Initialize();

            UpdatesAfter.Add(typeof(NodeGroupSystem));

            SubscribeLocalEvent<ApcPowerReceiverComponent, ComponentInit>(ApcPowerReceiverInit);
            SubscribeLocalEvent<ApcPowerReceiverComponent, ComponentShutdown>(ApcPowerReceiverShutdown);
            SubscribeLocalEvent<ApcPowerReceiverComponent, EntityPausedEvent>(ApcPowerReceiverPaused);
            SubscribeLocalEvent<PowerNetworkBatteryComponent, ComponentInit>(BatteryInit);
            SubscribeLocalEvent<PowerNetworkBatteryComponent, ComponentShutdown>(BatteryShutdown);
            SubscribeLocalEvent<PowerNetworkBatteryComponent, EntityPausedEvent>(BatteryPaused);
            SubscribeLocalEvent<PowerConsumerComponent, ComponentInit>(PowerConsumerInit);
            SubscribeLocalEvent<PowerConsumerComponent, ComponentShutdown>(PowerConsumerShutdown);
            SubscribeLocalEvent<PowerConsumerComponent, EntityPausedEvent>(PowerConsumerPaused);
            SubscribeLocalEvent<PowerSupplierComponent, ComponentInit>(PowerSupplierInit);
            SubscribeLocalEvent<PowerSupplierComponent, ComponentShutdown>(PowerSupplierShutdown);
            SubscribeLocalEvent<PowerSupplierComponent, EntityPausedEvent>(PowerSupplierPaused);
        }

        private void ApcPowerReceiverInit(EntityUid uid, ApcPowerReceiverComponent component, ComponentInit args)
        {
            AllocLoad(component.NetworkLoad);
        }

        private void ApcPowerReceiverShutdown(EntityUid uid, ApcPowerReceiverComponent component,
            ComponentShutdown args)
        {
            _powerState.Loads.Remove(component.NetworkLoad.Id);
        }

        private static void ApcPowerReceiverPaused(
            EntityUid uid,
            ApcPowerReceiverComponent component,
            EntityPausedEvent args)
        {
            component.NetworkLoad.Paused = args.Paused;
        }

        private void BatteryInit(EntityUid uid, PowerNetworkBatteryComponent component, ComponentInit args)
        {
            AllocBattery(component.NetworkBattery);
        }

        private void BatteryShutdown(EntityUid uid, PowerNetworkBatteryComponent component, ComponentShutdown args)
        {
            _powerState.Batteries.Remove(component.NetworkBattery.Id);
        }

        private static void BatteryPaused(EntityUid uid, PowerNetworkBatteryComponent component, EntityPausedEvent args)
        {
            component.NetworkBattery.Paused = args.Paused;
        }

        private void PowerConsumerInit(EntityUid uid, PowerConsumerComponent component, ComponentInit args)
        {
            AllocLoad(component.NetworkLoad);
        }

        private void PowerConsumerShutdown(EntityUid uid, PowerConsumerComponent component, ComponentShutdown args)
        {
            _powerState.Loads.Remove(component.NetworkLoad.Id);
        }

        private static void PowerConsumerPaused(EntityUid uid, PowerConsumerComponent component, EntityPausedEvent args)
        {
            component.NetworkLoad.Paused = args.Paused;
        }

        private void PowerSupplierInit(EntityUid uid, PowerSupplierComponent component, ComponentInit args)
        {
            AllocSupply(component.NetworkSupply);
        }

        private void PowerSupplierShutdown(EntityUid uid, PowerSupplierComponent component, ComponentShutdown args)
        {
            _powerState.Supplies.Remove(component.NetworkSupply.Id);
        }

        private static void PowerSupplierPaused(EntityUid uid, PowerSupplierComponent component, EntityPausedEvent args)
        {
            component.NetworkSupply.Paused = args.Paused;
        }

        public void InitPowerNet(PowerNet powerNet)
        {
            AllocNetwork(powerNet.NetworkNode);
        }

        public void DestroyPowerNet(PowerNet powerNet)
        {
            _powerState.Networks.Remove(powerNet.NetworkNode.Id);
        }

        public void QueueReconnectPowerNet(PowerNet powerNet)
        {
            _powerNetReconnectQueue.Add(powerNet);
        }

        public void InitApcNet(ApcNet apcNet)
        {
            AllocNetwork(apcNet.NetworkNode);
        }

        public void DestroyApcNet(ApcNet apcNet)
        {
            _powerState.Networks.Remove(apcNet.NetworkNode.Id);
        }

        public void QueueReconnectApcNet(ApcNet apcNet)
        {
            _apcNetReconnectQueue.Add(apcNet);
        }

        public PowerStatistics GetStatistics()
        {
            return new()
            {
                CountBatteries = _powerState.Batteries.Count,
                CountLoads = _powerState.Loads.Count,
                CountNetworks = _powerState.Networks.Count,
                CountSupplies = _powerState.Supplies.Count
            };
        }

        public override void Update(float frameTime)
        {
            base.Update(frameTime);

            // Reconnect networks.
            {
                foreach (var apcNet in _apcNetReconnectQueue)
                {
                    if (apcNet.Removed)
                        continue;

                    DoReconnectApcNet(apcNet);
                }

                _apcNetReconnectQueue.Clear();

                foreach (var powerNet in _powerNetReconnectQueue)
                {
                    if (powerNet.Removed)
                        continue;

                    DoReconnectPowerNet(powerNet);
                }

                _powerNetReconnectQueue.Clear();
            }

            // Synchronize batteries
            foreach (var battery in ComponentManager.EntityQuery<PowerNetworkBatteryComponent>())
            {
                RaiseLocalEvent(battery.Owner.Uid, new NetworkBatteryPreSync());
            }

            // Run power solver.
            _solver.Tick(frameTime, _powerState);

            // Synchronize batteries, the other way around.
            foreach (var battery in ComponentManager.EntityQuery<PowerNetworkBatteryComponent>())
            {
                RaiseLocalEvent(battery.Owner.Uid, new NetworkBatteryPostSync());
            }

            // Send events where necessary.
            {
                foreach (var apcReceiver in ComponentManager.EntityQuery<ApcPowerReceiverComponent>())
                {
                    var recv = apcReceiver.NetworkLoad.ReceivingPower;
                    ref var last = ref apcReceiver.LastPowerReceived;

                    if (!MathHelper.CloseTo(recv, last))
                    {
                        last = recv;
                        apcReceiver.ApcPowerChanged();
                    }
                }

                foreach (var consumer in ComponentManager.EntityQuery<PowerConsumerComponent>())
                {
                    var newRecv = consumer.NetworkLoad.ReceivingPower;
                    ref var lastRecv = ref consumer.LastReceived;
                    if (!MathHelper.CloseTo(lastRecv, newRecv))
                    {
                        lastRecv = newRecv;
                        var msg = new PowerConsumerReceivedChanged(newRecv, consumer.DrawRate);
                        RaiseLocalEvent(consumer.Owner.Uid, msg);
                    }
                }
            }
        }

        private void AllocLoad(PowerState.Load load)
        {
            load.Id = AllocId();
            _powerState.Loads.Add(load.Id, load);
        }

        private void AllocSupply(PowerState.Supply supply)
        {
            supply.Id = AllocId();
            _powerState.Supplies.Add(supply.Id, supply);
        }

        private void AllocBattery(PowerState.Battery battery)
        {
            battery.Id = AllocId();
            _powerState.Batteries.Add(battery.Id, battery);
        }

        private void AllocNetwork(PowerState.Network network)
        {
            network.Id = AllocId();
            _powerState.Networks.Add(network.Id, network);
        }

        private static void DoReconnectApcNet(ApcNet net)
        {
            var netNode = net.NetworkNode;

            netNode.Loads.Clear();
            netNode.BatteriesDischarging.Clear();
            netNode.BatteriesCharging.Clear();
            netNode.Supplies.Clear();

            foreach (var provider in net.Providers)
            {
                foreach (var receiver in provider.LinkedReceivers)
                {
                    netNode.Loads.Add(receiver.NetworkLoad.Id);
                    receiver.NetworkLoad.LinkedNetwork = netNode.Id;
                }
            }

            foreach (var apc in net.Apcs)
            {
                var netBattery = apc.Owner.GetComponent<PowerNetworkBatteryComponent>();
                netNode.BatteriesDischarging.Add(netBattery.NetworkBattery.Id);
                netBattery.NetworkBattery.LinkedNetworkDischarging = netNode.Id;
            }
        }

        private static void DoReconnectPowerNet(PowerNet net)
        {
            var netNode = net.NetworkNode;

            netNode.Loads.Clear();
            netNode.Supplies.Clear();
            netNode.BatteriesCharging.Clear();
            netNode.BatteriesDischarging.Clear();

            foreach (var consumer in net.Consumers)
            {
                netNode.Loads.Add(consumer.NetworkLoad.Id);
                consumer.NetworkLoad.LinkedNetwork = netNode.Id;
            }

            foreach (var supplier in net.Suppliers)
            {
                netNode.Supplies.Add(supplier.NetworkSupply.Id);
                supplier.NetworkSupply.LinkedNetwork = netNode.Id;
            }

            foreach (var charger in net.Chargers)
            {
                var battery = charger.Owner.GetComponent<PowerNetworkBatteryComponent>();
                netNode.BatteriesCharging.Add(battery.NetworkBattery.Id);
                battery.NetworkBattery.LinkedNetworkCharging = netNode.Id;
            }

            foreach (var discharger in net.Dischargers)
            {
                var battery = discharger.Owner.GetComponent<PowerNetworkBatteryComponent>();
                netNode.BatteriesDischarging.Add(battery.NetworkBattery.Id);
                battery.NetworkBattery.LinkedNetworkDischarging = netNode.Id;
            }
        }

        private PowerState.NodeId AllocId()
        {
            return new(_nextId++);
        }
    }

    /// <summary>
    ///     Raised before power network simulation happens, to synchronize battery state from
    ///     components like <see cref="BatteryComponent"/> into <see cref="PowerNetworkBatteryComponent"/>.
    /// </summary>
    public sealed class NetworkBatteryPreSync : EntityEventArgs
    {
    }

    /// <summary>
    ///     Raised after power network simulation happens, to synchronize battery charge changes from
    ///     <see cref="PowerNetworkBatteryComponent"/> to components like <see cref="BatteryComponent"/>.
    /// </summary>
    public sealed class NetworkBatteryPostSync : EntityEventArgs
    {
    }

    /// <summary>
    ///     Raised when the amount of receiving power on a <see cref="PowerConsumerComponent"/> changes.
    /// </summary>
    public sealed class PowerConsumerReceivedChanged : EntityEventArgs
    {
        public float ReceivedPower { get; }
        public float DrawRate { get; }

        public PowerConsumerReceivedChanged(float receivedPower, float drawRate)
        {
            ReceivedPower = receivedPower;
            DrawRate = drawRate;
        }
    }

    public struct PowerStatistics
    {
        public int CountNetworks;
        public int CountLoads;
        public int CountSupplies;
        public int CountBatteries;
    }
}