From cfab87b3d713249aca4c411b0739f15871ad421d Mon Sep 17 00:00:00 2001
From: Nemanja <98561806+EmoGarbage404@users.noreply.github.com>
Date: Sat, 22 Oct 2022 18:38:57 -0400
Subject: [PATCH] generator machine upgrading (#12145)

---
 .../Components/MachineComponent.cs            |  9 +++++
 .../Components/UpgradePowerDrawComponent.cs   | 14 ++-----
 .../UpgradePowerSupplierComponent.cs          | 30 ++++++++++++++
 .../EntitySystems/UpgradeBatterySystem.cs     |  2 +-
 ...werDrawSystem.cs => UpgradePowerSystem.cs} | 39 +++++++++++++++++--
 .../Power/Generation/generators.yml           |  6 +++
 6 files changed, 85 insertions(+), 15 deletions(-)
 create mode 100644 Content.Server/Power/Components/UpgradePowerSupplierComponent.cs
 rename Content.Server/Power/EntitySystems/{UpgradePowerDrawSystem.cs => UpgradePowerSystem.cs} (50%)

diff --git a/Content.Server/Construction/Components/MachineComponent.cs b/Content.Server/Construction/Components/MachineComponent.cs
index 82716e9067..4e3013d1b1 100644
--- a/Content.Server/Construction/Components/MachineComponent.cs
+++ b/Content.Server/Construction/Components/MachineComponent.cs
@@ -13,4 +13,13 @@ namespace Content.Server.Construction.Components
         public Container BoardContainer = default!;
         public Container PartContainer = default!;
     }
+
+    /// <summary>
+    /// The different types of scaling that are available for machine upgrades
+    /// </summary>
+    public enum MachineUpgradeScalingType : byte
+    {
+        Linear,
+        Exponential
+    }
 }
diff --git a/Content.Server/Power/Components/UpgradePowerDrawComponent.cs b/Content.Server/Power/Components/UpgradePowerDrawComponent.cs
index 6fbd1e97f1..2ecd1d4a05 100644
--- a/Content.Server/Power/Components/UpgradePowerDrawComponent.cs
+++ b/Content.Server/Power/Components/UpgradePowerDrawComponent.cs
@@ -1,4 +1,5 @@
-using Content.Shared.Construction.Prototypes;
+using Content.Server.Construction.Components;
+using Content.Shared.Construction.Prototypes;
 using Robust.Shared.Serialization.TypeSerializers.Implementations.Custom.Prototype;
 
 namespace Content.Server.Power.Components;
@@ -34,14 +35,7 @@ public sealed class UpgradePowerDrawComponent : Component
     /// What type of scaling is being used?
     /// </summary>
     [DataField("scaling", required: true), ViewVariables(VVAccess.ReadWrite)]
-    public PowerDrawScalingType Scaling;
+    public MachineUpgradeScalingType Scaling;
 }
 
-/// <summary>
-/// The different types of scaling that are available
-/// </summary>
-public enum PowerDrawScalingType : byte
-{
-    Linear,
-    Exponential
-}
+
diff --git a/Content.Server/Power/Components/UpgradePowerSupplierComponent.cs b/Content.Server/Power/Components/UpgradePowerSupplierComponent.cs
new file mode 100644
index 0000000000..25cf38155f
--- /dev/null
+++ b/Content.Server/Power/Components/UpgradePowerSupplierComponent.cs
@@ -0,0 +1,30 @@
+using Content.Server.Construction.Components;
+using Content.Shared.Construction.Prototypes;
+using Robust.Shared.Serialization.TypeSerializers.Implementations.Custom.Prototype;
+
+namespace Content.Server.Power.Components;
+
+[RegisterComponent]
+public sealed class UpgradePowerSupplierComponent : Component
+{
+    [ViewVariables(VVAccess.ReadWrite)]
+    public float BaseSupplyRate;
+
+    /// <summary>
+    /// The machine part that affects the power supplu.
+    /// </summary>
+    [DataField("machinePartPowerSupply", customTypeSerializer: typeof(PrototypeIdSerializer<MachinePartPrototype>)), ViewVariables(VVAccess.ReadWrite)]
+    public string MachinePartPowerSupply = "Capacitor";
+
+    /// <summary>
+    /// The multiplier used for scaling the power supply.
+    /// </summary>
+    [DataField("powerSupplyMultiplier", required: true), ViewVariables(VVAccess.ReadWrite)]
+    public float PowerSupplyMultiplier = 1f;
+
+    /// <summary>
+    /// What type of scaling is being used?
+    /// </summary>
+    [DataField("scaling", required: true), ViewVariables(VVAccess.ReadWrite)]
+    public MachineUpgradeScalingType Scaling;
+}
diff --git a/Content.Server/Power/EntitySystems/UpgradeBatterySystem.cs b/Content.Server/Power/EntitySystems/UpgradeBatterySystem.cs
index 0d48b76463..55698a7ff6 100644
--- a/Content.Server/Power/EntitySystems/UpgradeBatterySystem.cs
+++ b/Content.Server/Power/EntitySystems/UpgradeBatterySystem.cs
@@ -2,7 +2,7 @@ using Content.Server.Construction;
 using Content.Server.Power.Components;
 using JetBrains.Annotations;
 
-namespace Content.Server.Power.SMES
+namespace Content.Server.Power.EntitySystems
 {
     [UsedImplicitly]
     public sealed class UpgradeBatterySystem : EntitySystem
diff --git a/Content.Server/Power/EntitySystems/UpgradePowerDrawSystem.cs b/Content.Server/Power/EntitySystems/UpgradePowerSystem.cs
similarity index 50%
rename from Content.Server/Power/EntitySystems/UpgradePowerDrawSystem.cs
rename to Content.Server/Power/EntitySystems/UpgradePowerSystem.cs
index 62aa508059..f20b330bbb 100644
--- a/Content.Server/Power/EntitySystems/UpgradePowerDrawSystem.cs
+++ b/Content.Server/Power/EntitySystems/UpgradePowerSystem.cs
@@ -1,19 +1,23 @@
 using Content.Server.Construction;
+using Content.Server.Construction.Components;
 using Content.Server.Power.Components;
 
 namespace Content.Server.Power.EntitySystems;
 
 /// <summary>
 /// This handles using upgraded machine parts
-/// to modify the power load of a machine.
+/// to modify the power supply/generation of a machine.
 /// </summary>
-public sealed class UpgradePowerDrawSystem : EntitySystem
+public sealed class UpgradePowerSystem : EntitySystem
 {
     /// <inheritdoc/>
     public override void Initialize()
     {
         SubscribeLocalEvent<UpgradePowerDrawComponent, MapInitEvent>(OnMapInit);
         SubscribeLocalEvent<UpgradePowerDrawComponent, RefreshPartsEvent>(OnRefreshParts);
+
+        SubscribeLocalEvent<UpgradePowerSupplierComponent, MapInitEvent>(OnSupplierMapInit);
+        SubscribeLocalEvent<UpgradePowerSupplierComponent, RefreshPartsEvent>(OnSupplierRefreshParts);
     }
 
     private void OnMapInit(EntityUid uid, UpgradePowerDrawComponent component, MapInitEvent args)
@@ -30,10 +34,10 @@ public sealed class UpgradePowerDrawSystem : EntitySystem
         var rating = args.PartRatings[component.MachinePartPowerDraw];
         switch (component.Scaling)
         {
-            case PowerDrawScalingType.Linear:
+            case MachineUpgradeScalingType.Linear:
                 load += component.PowerDrawMultiplier * (rating - 1);
                 break;
-            case PowerDrawScalingType.Exponential:
+            case MachineUpgradeScalingType.Exponential:
                 load *= MathF.Pow(component.PowerDrawMultiplier, rating - 1);
                 break;
             default:
@@ -46,4 +50,31 @@ public sealed class UpgradePowerDrawSystem : EntitySystem
         if (TryComp<PowerConsumerComponent>(uid, out var powa2))
             powa2.DrawRate = load;
     }
+
+    private void OnSupplierMapInit(EntityUid uid, UpgradePowerSupplierComponent component, MapInitEvent args)
+    {
+        if (TryComp<PowerSupplierComponent>(uid, out var supplier))
+            component.BaseSupplyRate = supplier.MaxSupply;
+    }
+
+    private void OnSupplierRefreshParts(EntityUid uid, UpgradePowerSupplierComponent component, RefreshPartsEvent args)
+    {
+        if (!TryComp<PowerSupplierComponent>(uid, out var powa))
+            return;
+
+        var rating = args.PartRatings[component.MachinePartPowerSupply];
+        switch (component.Scaling)
+        {
+            case MachineUpgradeScalingType.Linear:
+                powa.MaxSupply += component.BaseSupplyRate * (rating - 1);
+                break;
+            case MachineUpgradeScalingType.Exponential:
+                powa.MaxSupply *= MathF.Pow(component.PowerSupplyMultiplier, rating - 1);
+                break;
+            default:
+                Logger.Error($"invalid power scaling type for {ToPrettyString(uid)}.");
+                powa.MaxSupply = component.BaseSupplyRate;
+                break;
+        }
+    }
 }
diff --git a/Resources/Prototypes/Entities/Structures/Power/Generation/generators.yml b/Resources/Prototypes/Entities/Structures/Power/Generation/generators.yml
index cc1d6e2fce..4ba2270553 100644
--- a/Resources/Prototypes/Entities/Structures/Power/Generation/generators.yml
+++ b/Resources/Prototypes/Entities/Structures/Power/Generation/generators.yml
@@ -169,6 +169,9 @@
     LayoutId: GeneratorPlasma
   - type: Machine
     board: GeneratorPlasmaMachineCircuitboard
+  - type: UpgradePowerSupplier
+    powerSupplyMultiplier: 1.25
+    scaling: Exponential
 
 - type: entity
   parent: [ BaseGenerator, ConstructibleMachine ]
@@ -185,6 +188,9 @@
     LayoutId: GeneratorUranium
   - type: Machine
     board: GeneratorUraniumMachineCircuitboard
+  - type: UpgradePowerSupplier
+    powerSupplyMultiplier: 1.25
+    scaling: Exponential
 
 - type: entity
   parent: BaseGeneratorWallmount