Fixes secborg batons (#16196)

* Cell linking

* V2

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>")

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)

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)

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"