From 6b3e56e754c11df67587b6ec78741785fd3c3c17 Mon Sep 17 00:00:00 2001
From: Jbleezy <thejbirdjustin@yahoo.com>
Date: Sat, 4 Mar 2023 22:45:10 -0800
Subject: [PATCH] Staffs: full ammo when initially picked up after upgrading

---
 README.md                                  |  1 +
 scripts/zm/replaced/zm_tomb_main_quest.gsc | 43 ++++++++++++++++++++++
 scripts/zm/zm_tomb/zm_tomb_reimagined.gsc  |  1 +
 3 files changed, 45 insertions(+)

diff --git a/README.md b/README.md
index 484f0d22..ec033dd8 100644
--- a/README.md
+++ b/README.md
@@ -199,6 +199,7 @@
 * Upgrade only requires collecting souls in The Crazy Place
 * Can no longer collect souls in The Crazy Place from staff kills
 * Can collect souls in The Crazy Place after upgrading to fill 5% of max ammo per soul
+* Get full ammo when initially picked up after upgrading (normally missing one clip)
 * Max ammo no longer fills clip ammo
 
 ## Wallbuys
diff --git a/scripts/zm/replaced/zm_tomb_main_quest.gsc b/scripts/zm/replaced/zm_tomb_main_quest.gsc
index e3396ec1..f5f969b3 100644
--- a/scripts/zm/replaced/zm_tomb_main_quest.gsc
+++ b/scripts/zm/replaced/zm_tomb_main_quest.gsc
@@ -253,6 +253,49 @@ waittill_staff_inserted()
     }
 }
 
+watch_for_player_pickup_staff()
+{
+    staff_picked_up = 0;
+    pickup_message = self staff_get_pickup_message();
+    self.trigger set_unitrigger_hint_string( pickup_message );
+    self show();
+    self.trigger trigger_on();
+
+    while ( !staff_picked_up )
+    {
+        self.trigger waittill( "trigger", player );
+
+        self notify( "retrieved", player );
+
+        if ( player can_pickup_staff() )
+        {
+            weapon_drop = player getcurrentweapon();
+            a_weapons = player getweaponslistprimaries();
+            n_max_other_weapons = get_player_weapon_limit( player ) - 1;
+
+            if ( a_weapons.size > n_max_other_weapons || issubstr( weapon_drop, "staff" ) )
+                player takeweapon( weapon_drop );
+
+            player thread watch_staff_ammo_reload();
+            self ghost();
+            self setinvisibletoall();
+            player giveweapon( self.weapname );
+            player switchtoweapon( self.weapname );
+            clip_size = weaponclipsize( self.weapname );
+            player setweaponammoclip( self.weapname, clip_size );
+            player givemaxammo( self.weapname );
+            self.owner = player;
+            level notify( "stop_staff_sound" );
+            self notify( "staff_equip" );
+            staff_picked_up = 1;
+            self.charger.is_inserted = 0;
+            self setclientfield( "staff_charger", 0 );
+            self.charger.full = 1;
+            maps\mp\zm_tomb_craftables::set_player_staff( self.weapname, player );
+        }
+    }
+}
+
 staff_upgraded_reload()
 {
     self endon( "staff_equip" );
diff --git a/scripts/zm/zm_tomb/zm_tomb_reimagined.gsc b/scripts/zm/zm_tomb/zm_tomb_reimagined.gsc
index a9f7709f..34dacae3 100644
--- a/scripts/zm/zm_tomb/zm_tomb_reimagined.gsc
+++ b/scripts/zm/zm_tomb/zm_tomb_reimagined.gsc
@@ -14,6 +14,7 @@
 main()
 {
 	replaceFunc(maps\mp\zm_tomb_main_quest::main_quest_init, scripts\zm\replaced\zm_tomb_main_quest::main_quest_init);
+	replaceFunc(maps\mp\zm_tomb_main_quest::watch_for_player_pickup_staff, scripts\zm\replaced\zm_tomb_main_quest::watch_for_player_pickup_staff);
 	replaceFunc(maps\mp\zm_tomb_main_quest::staff_upgraded_reload, scripts\zm\replaced\zm_tomb_main_quest::staff_upgraded_reload);
 	replaceFunc(maps\mp\zm_tomb_main_quest::watch_staff_ammo_reload, scripts\zm\replaced\zm_tomb_main_quest::watch_staff_ammo_reload);
 	replaceFunc(maps\mp\zm_tomb_quest_air::air_puzzle_1_run, scripts\zm\replaced\zm_tomb_quest_air::air_puzzle_1_run);