From dd0243509551c14b7ada6949c06b7b9d5bff4277 Mon Sep 17 00:00:00 2001
From: SabreML <57483089+SabreML@users.noreply.github.com>
Date: Sat, 19 Jun 2021 14:10:01 +0100
Subject: [PATCH] Fixes secborg batons (#16196)

* Cell linking

* V2
---
 code/game/objects/items/weapons/stunbaton.dm  | 26 ++++++++++++++-----
 .../modules/mob/living/silicon/robot/robot.dm |  1 +
 .../mob/living/silicon/robot/robot_defense.dm |  1 +
 .../mob/living/silicon/robot/robot_modules.dm | 15 +++++++++++
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/code/game/objects/items/weapons/stunbaton.dm b/code/game/objects/items/weapons/stunbaton.dm
index 0df5817f598..64cab3bef7a 100644
--- a/code/game/objects/items/weapons/stunbaton.dm
+++ b/code/game/objects/items/weapons/stunbaton.dm
@@ -25,16 +25,30 @@
 	update_icon()
 
 /obj/item/melee/baton/loaded/Initialize(mapload) //this one starts with a cell pre-installed.
-	if(isrobot(loc.loc)) // First loc would be the module
+	link_new_cell()
+	return ..()
+
+/obj/item/melee/baton/Destroy()
+	if(cell?.loc == src)
+		QDEL_NULL(cell)
+	return ..()
+
+/**
+ * Updates the linked power cell on the baton.
+ *
+ * If the baton is held by a cyborg, link it to their internal cell.
+ * Else, spawn a new cell and use that instead.
+ * Arguments:
+ * * unlink - If TRUE, sets the `cell` variable to `null` rather than linking it to a new one.
+ */
+/obj/item/melee/baton/proc/link_new_cell(unlink = FALSE)
+	if(unlink)
+		cell = null
+	else if(isrobot(loc.loc)) // First loc is the module
 		var/mob/living/silicon/robot/R = loc.loc
 		cell = R.cell
 	else
 		cell = new(src)
-	return ..()
-
-/obj/item/melee/baton/Destroy()
-	QDEL_NULL(cell)
-	return ..()
 
 /obj/item/melee/baton/suicide_act(mob/user)
 	user.visible_message("<span class='suicide'>[user] is putting the live [name] in [user.p_their()] mouth! It looks like [user.p_theyre()] trying to commit suicide.</span>")
diff --git a/code/modules/mob/living/silicon/robot/robot.dm b/code/modules/mob/living/silicon/robot/robot.dm
index 9bacb40fff2..f8c27db8dc7 100644
--- a/code/modules/mob/living/silicon/robot/robot.dm
+++ b/code/modules/mob/living/silicon/robot/robot.dm
@@ -757,6 +757,7 @@ GLOBAL_LIST_INIT(robot_verbs_default, list(
 			//This will mean that removing and replacing a power cell will repair the mount, but I don't care at this point. ~Z
 			C.brute_damage = 0
 			C.electronics_damage = 0
+			module?.update_cells()
 			diag_hud_set_borgcell()
 
 	else if(istype(W, /obj/item/encryptionkey/) && opened)
diff --git a/code/modules/mob/living/silicon/robot/robot_defense.dm b/code/modules/mob/living/silicon/robot/robot_defense.dm
index 84eb782833b..6469bfbf6a6 100644
--- a/code/modules/mob/living/silicon/robot/robot_defense.dm
+++ b/code/modules/mob/living/silicon/robot/robot_defense.dm
@@ -47,6 +47,7 @@
 			var/datum/robot_component/C = components["power cell"]
 			C.installed = 0
 			C.uninstall()
+			module?.update_cells(TRUE)
 			diag_hud_set_borgcell()
 
 	if(!opened)
diff --git a/code/modules/mob/living/silicon/robot/robot_modules.dm b/code/modules/mob/living/silicon/robot/robot_modules.dm
index 26d15986d86..9e944853a0f 100644
--- a/code/modules/mob/living/silicon/robot/robot_modules.dm
+++ b/code/modules/mob/living/silicon/robot/robot_modules.dm
@@ -188,6 +188,16 @@
 		var/obj/item/I = item
 		I.cyborg_recharge(coeff, R.emagged)
 
+/**
+ * Called when the robot owner of this module has their power cell replaced.
+ *
+ * Changes the linked power cell for module items to the newly inserted cell, or to `null`.
+ * Arguments:
+ * * unlink_cell - If TRUE, set the item's power cell variable to `null` rather than linking it to a new one.
+ */
+/obj/item/robot_module/proc/update_cells(unlink_cell = FALSE)
+	return
+
 /**
  * Called when the robot owner of this module has the `unemag()` proc called on them, which is only via admin means.
  *
@@ -408,6 +418,11 @@
 	emag_modules = list(/obj/item/gun/energy/laser/cyborg)
 	special_rechargables = list(/obj/item/melee/baton/loaded, /obj/item/gun/energy/disabler/cyborg)
 
+/obj/item/robot_module/security/update_cells(unlink_cell = FALSE)
+	var/obj/item/melee/baton/B = locate(/obj/item/melee/baton/loaded) in modules
+	if(B)
+		B.link_new_cell(unlink_cell)
+
 // Janitor cyborg module.
 /obj/item/robot_module/janitor
 	name = "janitorial robot module"