Implement energy regen patch

Projects/TerraTech/TerraTech/Class1.cs

@@ -11,6 +11,7 @@ using UnityEngine;
 // TODO: Make more energy generate
 // TODO: Make attractor more range
 // TODO: Make battery bigger
+// TODO: Make shield and repair bigger
 
 namespace TerraTech {
     [BepInPlugin(pluginGuid, pluginName, pluginVersion)]
@@ -35,7 +36,7 @@ namespace TerraTech {
             allProjectilesHoming = Config.Bind("General", "Make All Projectiles Home", false);
 
             shootingSpeedMultiplier.SettingChanged += (sender, args) => WeaponPropertiesManager.DoPatch();
-            WeaponPropertiesManager.DoPatch();
+            energyGenMultiplier.SettingChanged += (sender, args) => GeneratorPropertiesManager.DoPatch();
 
             Logger.LogInfo("Cyka mod loaded");
             HarmonyFileLog.Enabled = true;
@@ -59,15 +60,6 @@ namespace TerraTech {
         }
     }
 
-    [HarmonyPatch(typeof(ModuleEnergy), "Supply")]
-    public class EnergyGenerationMultiplierPatch {
-        static void Prefix(EnergyRegulator.EnergyType type, ref float energyAmount) {
-            if (energyAmount > 0) {
-                energyAmount *= Main.energyGenMultiplier.Value;
-            }
-        }
-    }
-
     [HarmonyPatch(typeof(Projectile), "Fire")]
     public class HomingProjectilePatch {
         private static Dictionary<FireData, bool> patchedObjects = new Dictionary<FireData, bool>();
@@ -88,57 +80,34 @@ namespace TerraTech {
         private static Dictionary<ModuleWeaponGun, float> weapons = new Dictionary<ModuleWeaponGun, float>();
 
         [HarmonyPostfix]
-        [HarmonyPatch(typeof(ModuleWeaponGun), "OnPool")]
-        static void Postfix1WeaponCreate(ModuleWeaponGun __instance) {
-            AddIfMissing(__instance);
-        }
-
-        [HarmonyPostfix]
-        [HarmonyPatch(typeof(ModuleWeaponGun), "OnSpawn")]
-        static void Postfix2WeaponCreate(ModuleWeaponGun __instance) {
-            AddIfMissing(__instance);
-        }
-
-        private static void AddIfMissing(ModuleWeaponGun weaponGun) {
-            if (!weapons.ContainsKey(weaponGun)) {
-                weapons.Add(weaponGun, weaponGun.m_ShotCooldown);
-                if (ShouldPatch()) {
-                    DoPatchSingle(weaponGun);
-                }
+        [HarmonyPatch(typeof(ModuleWeaponGun), "OnAttached")]
+        static void PostfixCreate(ModuleWeaponGun __instance) {
+            if (!weapons.ContainsKey(__instance)) {
+                weapons.Add(__instance, __instance.m_ShotCooldown);
+                DoPatchSingle(__instance);
             }
         }
 
         [HarmonyPostfix]
-        [HarmonyPatch(typeof(ModuleWeaponGun), "OnRecycle")]
-        static void Postfix1WeaponDestroy(ModuleWeaponGun __instance) {
+        [HarmonyPatch(typeof(ModuleWeaponGun), "OnDetaching")]
+        static void PostfixDestroy(ModuleWeaponGun __instance) {
+            DoRestoreSingle(__instance);
             weapons.Remove(__instance);
         }
 
-        // [HarmonyPrefix]
-        // [HarmonyPatch(typeof(ModuleWeaponGun), "ProcessFiring")]
-        // static void PrefixProcessFiring(ModuleWeaponGun __instance, ref bool firing) {
-        //     if (firing) {
-        //         Console.WriteLine("" + cooldowns.ContainsKey(__instance) + " " + __instance.m_ShotCooldown);
-        //     }
-        // }
-
         public static void DoPatch() {
             Console.WriteLine("Modifying " + weapons.Count + " weapons");
-            Console.WriteLine("Should patch: " + ShouldPatch());
             foreach (KeyValuePair<ModuleWeaponGun, float> keyValuePair in weapons) {
-                if (ShouldPatch()) {
-                    DoPatchSingle(keyValuePair.Key);
-                } else {
-                    DoRestoreSingle(keyValuePair.Key);
-                }
+                DoRestoreSingle(keyValuePair.Key);
+                DoPatchSingle(keyValuePair.Key);
             }
         }
 
         static void DoPatchSingle(ModuleWeaponGun weapon) {
-            // Console.WriteLine("Patching " + weapon.name);
-            // Console.WriteLine("Old value " + weapon.m_ShotCooldown);
+            Console.WriteLine("Patching " + weapon.name);
+            Console.WriteLine("Old value " + weapon.m_ShotCooldown);
             weapon.m_ShotCooldown /= Main.shootingSpeedMultiplier.Value;
-            // Console.WriteLine("New value " + weapon.m_ShotCooldown);
+            Console.WriteLine("New value " + weapon.m_ShotCooldown);
         }
 
         static void DoRestoreSingle(ModuleWeaponGun weapon) {
@@ -146,9 +115,62 @@ namespace TerraTech {
                 weapon.m_ShotCooldown = weapons[weapon];
             }
         }
+    }
+
+    [HarmonyPatch]
+    public class GeneratorPropertiesManager {
+        private static Dictionary<ModuleEnergy, float> generators = new Dictionary<ModuleEnergy, float>();
+
+        [HarmonyPostfix]
+        [HarmonyPatch(typeof(ModuleEnergy), "OnAttached")]
+        static void PostfixCreate(ModuleEnergy __instance) {
+            if (!generators.ContainsKey(__instance)) {
+                generators.Add(__instance, GetValue(__instance));
+                if (ShouldPatch()) {
+                    DoPatchSingle(__instance);
+                }
+            }
+        }
+
+        [HarmonyPostfix]
+        [HarmonyPatch(typeof(ModuleEnergy), "OnDetaching")]
+        static void PostfixDestroy(ModuleEnergy __instance) {
+            DoRestoreSingle(__instance);
+            generators.Remove(__instance);
+        }
+
+        public static void DoPatch() {
+            Console.WriteLine("Modifying " + generators.Count + " generators");
+            Console.WriteLine("Should patch: " + ShouldPatch());
+            foreach (KeyValuePair<ModuleEnergy, float> keyValuePair in generators) {
+                DoRestoreSingle(keyValuePair.Key);
+                DoPatchSingle(keyValuePair.Key);
+            }
+        }
+
+        static void DoPatchSingle(ModuleEnergy moduleEnergy) {
+            Console.WriteLine("Patching " + moduleEnergy.name);
+            Console.WriteLine("Old value " + GetValue(moduleEnergy));
+            SetValue(moduleEnergy, GetValue(moduleEnergy) * Main.energyGenMultiplier.Value);
+            Console.WriteLine("New value " + GetValue(moduleEnergy));
+        }
+
+        static void DoRestoreSingle(ModuleEnergy moduleEnergy) {
+            if (generators.ContainsKey(moduleEnergy)) {
+                SetValue(moduleEnergy, generators[moduleEnergy]);
+            }
+        }
 
         static bool ShouldPatch() {
-            return Math.Abs(Main.shootingSpeedMultiplier.Value - 1f) > 0.0001f;
+            return Math.Abs(Main.energyGenMultiplier.Value - 1f) > 0.0001f;
+        }
+
+        private static float GetValue(ModuleEnergy moduleEnergy) {
+            return Traverse.Create(moduleEnergy).Field("m_OutputPerSecond").GetValue() as float? ?? 0f;
+        }
+
+        private static void SetValue(ModuleEnergy moduleEnergy, float value) {
+            Traverse.Create(moduleEnergy).Field("m_OutputPerSecond").SetValue(value);
         }
     }
 }