From 10adfcfe28e4c8aac21badd3051848d4233a8fb5 Mon Sep 17 00:00:00 2001
From: KJeff01 <kjeff1997@gmail.com>
Date: Thu, 28 Dec 2023 14:01:21 -0600
Subject: [PATCH] Prevent launch button disappearing on LZ compromise

A little regression from 452e21bc50b79b9d854b974e4256d20ad96b56c6.
---
 .../campaign/libcampaign_includes/victory.js  |  4 +--
 src/quickjs_backend.cpp                       |  2 +-
 src/wzapi.cpp                                 | 32 ++++++++++---------
 src/wzapi.h                                   |  2 +-
 4 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/data/base/script/campaign/libcampaign_includes/victory.js b/data/base/script/campaign/libcampaign_includes/victory.js
index 677ed0f6a8f..3029697c729 100644
--- a/data/base/script/campaign/libcampaign_includes/victory.js
+++ b/data/base/script/campaign/libcampaign_includes/victory.js
@@ -452,7 +452,7 @@ function __camVictoryOffworld()
 		//Protect against early access to reinforcements GUI if it shouldn't be available yet
 		if (__camVictoryData.reinforcements >= 0)
 		{
-			setReinforcementTime(LZ_COMPROMISED_TIME);
+			setReinforcementTime(LZ_COMPROMISED_TIME, false);
 		}
 		if (__camLZCompromisedTicker === 0)
 		{
@@ -474,7 +474,7 @@ function __camVictoryOffworld()
 		camTrace("LZ clear");
 		const pos = camMakePos(lz);
 		playSound(cam_sounds.lz.LZClear, pos.x, pos.y, 0);
-		setReinforcementTime(__camVictoryData.reinforcements);
+		setReinforcementTime(__camVictoryData.reinforcements, false);
 		__camLZCompromisedTicker = 0;
 		if (__camRTLZTicker === 0)
 		{
diff --git a/src/quickjs_backend.cpp b/src/quickjs_backend.cpp
index ab49257f1c8..ea94be5a906 100644
--- a/src/quickjs_backend.cpp
+++ b/src/quickjs_backend.cpp
@@ -3488,7 +3488,7 @@ bool quickjs_scripting_instance::registerFunctions(const std::string& scriptName
 	JS_REGISTER_FUNC(applyLimitSet, 0); // WZAPI
 	JS_REGISTER_FUNC(setMissionTime, 1); // WZAPI
 	JS_REGISTER_FUNC(getMissionTime, 0); // WZAPI
-	JS_REGISTER_FUNC(setReinforcementTime, 1); // WZAPI
+	JS_REGISTER_FUNC2(setReinforcementTime, 1, 2); // WZAPI
 	JS_REGISTER_FUNC2(completeResearch, 1, 3); // WZAPI
 	JS_REGISTER_FUNC2(completeAllResearch, 0, 1); // WZAPI
 	JS_REGISTER_FUNC2(enableResearch, 1, 2); // WZAPI
diff --git a/src/wzapi.cpp b/src/wzapi.cpp
index 15b99aa5edd..ad2650f9468 100644
--- a/src/wzapi.cpp
+++ b/src/wzapi.cpp
@@ -2511,7 +2511,7 @@ int wzapi::getMissionTime(WZAPI_NO_PARAMS)
 	return (mission.time - (gameTime - mission.startTime)) / GAME_TICKS_PER_SEC;
 }
 
-//-- ## setReinforcementTime(time)
+//-- ## setReinforcementTime(time[, removeLaunch])
 //--
 //-- Set time for reinforcements to arrive. If time is negative, the reinforcement GUI
 //-- is removed and the timer stopped. Time is in seconds.
@@ -2519,30 +2519,32 @@ int wzapi::getMissionTime(WZAPI_NO_PARAMS)
 //-- is set to "--:--" and reinforcements are suppressed until this function is called
 //-- again with a regular time value.
 //--
-wzapi::no_return_value wzapi::setReinforcementTime(WZAPI_PARAMS(int _time))
+wzapi::no_return_value wzapi::setReinforcementTime(WZAPI_PARAMS(int _time, optional<bool> _removeLaunch))
 {
 	int time = _time * GAME_TICKS_PER_SEC;
-	SCRIPT_ASSERT({}, context, time == LZ_COMPROMISED_TIME || time < 60 * 60 * GAME_TICKS_PER_SEC,
-	              "The transport timer cannot be set to more than 1 hour!");
+	bool removeLaunch = _removeLaunch.value_or(true);
+	SCRIPT_ASSERT({}, context, time == LZ_COMPROMISED_TIME || time < 60 * 60 * GAME_TICKS_PER_SEC, "The transport timer cannot be set to more than 1 hour!");
 	SCRIPT_ASSERT({}, context, selectedPlayer < MAX_PLAYERS, "Invalid selectedPlayer for current client: %" PRIu32 "", selectedPlayer);
 	mission.ETA = time;
-	
 	if (time < 0)
 	{
 		intRemoveTransporterTimer();
 	}
-	/* Search for a transport that is idle; if we can't find any, remove the launch button
-	 * since there's no transport to launch */
-	auto droidIt = std::find_if(apsDroidLists[selectedPlayer].begin(), apsDroidLists[selectedPlayer].end(), [](DROID* d)
+	if (removeLaunch)
 	{
-		return isTransporter(d) && !transporterFlying(d);
-	});
-	DROID* psDroid = droidIt != apsDroidLists[selectedPlayer].end() ? *droidIt : nullptr;
+		/* Search for a transport that is idle; if we can't find any, remove the launch button
+		 * since there's no transport to launch */
+		auto droidIt = std::find_if(apsDroidLists[selectedPlayer].begin(), apsDroidLists[selectedPlayer].end(), [](DROID* d)
+		{
+			return isTransporter(d) && !transporterFlying(d);
+		});
+		DROID* psDroid = droidIt != apsDroidLists[selectedPlayer].end() ? *droidIt : nullptr;
 
-	// Didn't find an idle transporter, we can remove the launch button
-	if (psDroid == nullptr)
-	{
-		intRemoveTransporterLaunch();
+		// Didn't find an idle transporter, we can remove the launch button
+		if (psDroid == nullptr)
+		{
+			intRemoveTransporterLaunch();
+		}
 	}
 	resetMissionWidgets();
 	addTransporterTimerInterface();
diff --git a/src/wzapi.h b/src/wzapi.h
index 3f7e1dc7b3e..1c6d4ae5b52 100644
--- a/src/wzapi.h
+++ b/src/wzapi.h
@@ -1075,7 +1075,7 @@ namespace wzapi
 	bool applyLimitSet(WZAPI_NO_PARAMS);
 	no_return_value setMissionTime(WZAPI_PARAMS(int _time));
 	int getMissionTime(WZAPI_NO_PARAMS);
-	no_return_value setReinforcementTime(WZAPI_PARAMS(int _time));
+	no_return_value setReinforcementTime(WZAPI_PARAMS(int _time, optional<bool> _removeLaunch));
 	no_return_value completeResearch(WZAPI_PARAMS(std::string researchName, optional<int> _player, optional<bool> _forceResearch));
 	no_return_value completeAllResearch(WZAPI_PARAMS(optional<int> _player));
 	bool enableResearch(WZAPI_PARAMS(std::string researchName, optional<int> _player));