From 19b6e05f30c81d934c95159d0717b4e77a0d4bec Mon Sep 17 00:00:00 2001
From: Denis Belan <belan.denis@gmail.com>
Date: Wed, 15 Nov 2023 22:19:57 +0300
Subject: [PATCH] Update WorkOrderCosts with...

... possibility to use MechDef variables in formulae for Repair elements
---
 .../Patches/MechBayPanel_Patches.cs           | 41 +++++++++++++++++++
 .../MechLabPanel_OrderItemRepair_Patches.cs   | 22 ++++++++++
 .../WorkOrderCosts/WorkOrderCostsHandler.cs   |  5 ++-
 3 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 source/CustomComponents/WorkOrderCosts/Patches/MechBayPanel_Patches.cs
 create mode 100644 source/CustomComponents/WorkOrderCosts/Patches/MechLabPanel_OrderItemRepair_Patches.cs

diff --git a/source/CustomComponents/WorkOrderCosts/Patches/MechBayPanel_Patches.cs b/source/CustomComponents/WorkOrderCosts/Patches/MechBayPanel_Patches.cs
new file mode 100644
index 0000000..f05c8ae
--- /dev/null
+++ b/source/CustomComponents/WorkOrderCosts/Patches/MechBayPanel_Patches.cs
@@ -0,0 +1,41 @@
+using BattleTech.UI;
+
+namespace CustomComponents;
+
+[HarmonyPatch(typeof(MechBayPanel), "OnRepairMech")]
+public static class MechBayPanel_OnRepairMech_Patches
+{
+    public static void Prefix(ref bool __runOriginal, MechBayPanel __instance, MechBayMechUnitElement mechElement)
+    {
+        if (!__runOriginal)
+        {
+            return;
+        }
+
+        WorkOrderCostsHandler.Shared.ActiveMechDef = mechElement.MechDef;
+    }
+
+    public static void Postfix(MechBayPanel __instance)
+    {
+        WorkOrderCostsHandler.Shared.ActiveMechDef = null;
+    }
+}
+
+[HarmonyPatch(typeof(MechBayPanel), "OrderItemRepair")]
+public static class MechBayPanel_OrderItemRepair_Patches
+{
+    public static void Prefix(ref bool __runOriginal, MechBayPanel __instance, IMechLabDraggableItem item)
+    {
+        if (!__runOriginal)
+        {
+            return;
+        }
+        
+        WorkOrderCostsHandler.Shared.ActiveMechDef = item.MechDef;
+    }
+
+    public static void Postfix(MechBayPanel __instance)
+    {
+        WorkOrderCostsHandler.Shared.ActiveMechDef = null;
+    }
+}
\ No newline at end of file
diff --git a/source/CustomComponents/WorkOrderCosts/Patches/MechLabPanel_OrderItemRepair_Patches.cs b/source/CustomComponents/WorkOrderCosts/Patches/MechLabPanel_OrderItemRepair_Patches.cs
new file mode 100644
index 0000000..7833c65
--- /dev/null
+++ b/source/CustomComponents/WorkOrderCosts/Patches/MechLabPanel_OrderItemRepair_Patches.cs
@@ -0,0 +1,22 @@
+using BattleTech.UI;
+
+namespace CustomComponents;
+
+[HarmonyPatch(typeof(MechLabPanel), "OrderItemRepair")]
+public static class MechLabPanel_OrderItemRepair_Patches
+{
+    public static void Prefix(bool __runOriginal, MechLabPanel __instance)
+    {    
+        if (!__runOriginal)
+        {
+            return;
+        }
+
+        WorkOrderCostsHandler.Shared.ActiveMechDef = __instance.Sim.GetMechByID(__instance.baseWorkOrder.MechID);
+    }
+
+    public static void Postfix(MechLabPanel __instance)
+    {
+        WorkOrderCostsHandler.Shared.ActiveMechDef = null;
+    }
+}
\ No newline at end of file
diff --git a/source/CustomComponents/WorkOrderCosts/WorkOrderCostsHandler.cs b/source/CustomComponents/WorkOrderCosts/WorkOrderCostsHandler.cs
index d89faaa..27d04b7 100644
--- a/source/CustomComponents/WorkOrderCosts/WorkOrderCostsHandler.cs
+++ b/source/CustomComponents/WorkOrderCosts/WorkOrderCostsHandler.cs
@@ -6,6 +6,7 @@ internal class WorkOrderCostsHandler
 {
     public static readonly WorkOrderCostsHandler Shared = new();
 
+    public MechDef ActiveMechDef { get; set; } = null;
     public void ComponentInstallWorkOrder(MechDef mechDef, MechComponentRef mechComponent, ChassisLocations newLocation, WorkOrderEntry_InstallComponent result)
     {
 
@@ -42,11 +43,11 @@ public void ComponentRepairWorkOrder(MechComponentRef mechComponent, bool isOnMe
 
         if (mechComponent.DamageLevel == ComponentDamageLevel.Destroyed)
         {
-            ApplyCosts(result, workOrderCosts.RepairDestroyed, null);
+            ApplyCosts(result, workOrderCosts.RepairDestroyed, isOnMech ? ActiveMechDef : null);
         }
         else
         {
-            ApplyCosts(result, workOrderCosts.Repair, null);
+            ApplyCosts(result, workOrderCosts.Repair, isOnMech ? ActiveMechDef : null);
         }
     }