From 6b9ab6ef1224752c8de44abf0516fdb40a570c31 Mon Sep 17 00:00:00 2001
From: Ovahlord <dreadkiller@gmx.de>
Date: Wed, 27 Sep 2023 18:16:15 +0200
Subject: [PATCH] Core/Pets: fixed showing pet actions bars for Hunter and
 Warlocks when relogging

---
 src/server/game/Entities/Player/Player.cpp    | 85 +++++--------------
 src/server/game/Entities/Player/Player.h      |  9 +-
 src/server/game/Entities/Unit/Unit.cpp        |  2 +-
 src/server/game/Handlers/PetHandler.cpp       |  3 +-
 src/server/game/Server/Packets/PetPackets.cpp |  3 +
 src/server/scripts/Spells/spell_generic.cpp   | 44 ++++++++++
 src/server/scripts/Spells/spell_hunter.cpp    | 10 +--
 7 files changed, 76 insertions(+), 80 deletions(-)

diff --git a/src/server/game/Entities/Player/Player.cpp b/src/server/game/Entities/Player/Player.cpp
index 39545f95a0..5429e5d6dc 100644
--- a/src/server/game/Entities/Player/Player.cpp
+++ b/src/server/game/Entities/Player/Player.cpp
@@ -17875,30 +17875,6 @@ void Player::AbandonPet()
     _playerPetDataMap.erase(itr);
 }
 
-PlayerPetData* Player::GetPlayerPetDataCurrent()
-{
-    return nullptr;
-}
-
-Optional<uint8> Player::GetFirstUnusedActivePetSlot()
-{
-    std::set<uint8> unusedActiveSlot = { 0, 1, 2, 3, 4 }; //unfiltered
-    return Optional<uint8>{};
-}
-
-Optional<uint8> Player::GetFirstUnusedPetSlot()
-{
-    return Optional<uint8>{};
-}
-
-void Player::DeleteFromPlayerPetDataStore(uint32 petNumber)
-{
-}
-
-void Player::AddToPlayerPetDataStore(PlayerPetData&& playerPetData)
-{
-}
-
 void Player::_LoadQuestStatus(PreparedQueryResult result)
 {
     uint16 slot = 0;
@@ -20417,7 +20393,7 @@ bool Player::RemoveMItem(uint32 id)
     return mMitems.erase(id) ? true : false;
 }
 
-void Player::SendPetSpellsMessage(NewPet* pet)
+void Player::SendPetSpellsMessage(NewPet* pet, bool remove /*= false*/)
 {
     // Warlocks and Hunters cannot control their pets until they have learned their respective control spell
     if (pet->IsClassPet() && !CanControlClassPets())
@@ -20431,26 +20407,29 @@ void Player::SendPetSpellsMessage(NewPet* pet)
     }
 
     WorldPackets::Pet::PetSpellsMessage packet;
-    packet.PetGUID = pet->GetGUID();
-    packet._CreatureFamily = pet->GetCreatureTemplate()->family;         // creature family (required for pet talents)
-    packet.TimeLimit = pet->GetRemainingSummonDuration().count();
-    packet.ReactState = pet->GetReactState();
-    packet.CommandState = charmInfo->GetCommandState();
-    packet.Flag = 0;
-
-    charmInfo->BuildActionBar(packet.ActionButtons);
-    if (pet->IsClassPet())
+    if (!remove)
     {
-        for (auto const& itr : pet->GetSpells())
+        packet.PetGUID = pet->GetGUID();
+        packet._CreatureFamily = pet->GetCreatureTemplate()->family;         // creature family (required for pet talents)
+        packet.TimeLimit = pet->GetRemainingSummonDuration().count();
+        packet.ReactState = pet->GetReactState();
+        packet.CommandState = charmInfo->GetCommandState();
+        packet.Flag = 0;
+
+        charmInfo->BuildActionBar(packet.ActionButtons);
+        if (pet->IsClassPet())
         {
-            if (itr.second.State == PETSPELL_REMOVED)
-                continue;
+            for (auto const& itr : pet->GetSpells())
+            {
+                if (itr.second.State == PETSPELL_REMOVED)
+                    continue;
 
-            packet.Actions.push_back(MAKE_UNIT_ACTION_BUTTON(itr.first, itr.second.Active));
-        }
+                packet.Actions.push_back(MAKE_UNIT_ACTION_BUTTON(itr.first, itr.second.Active));
+            }
 
-        // Cooldowns
-        pet->GetSpellHistory()->WritePetSpellHistory(packet);
+            // Cooldowns
+            pet->GetSpellHistory()->WritePetSpellHistory(packet);
+        }
     }
 
     SendDirectMessage(packet.Write());
@@ -20476,30 +20455,6 @@ void Player::SendOnCancelExpectedVehicleRideAura() const
     SendDirectMessage(&data);
 }
 
-bool Player::CanControlPet(uint32 spellId) const
-{
-    // Hunters and Warlocks cannot control their pets until they learned their spell
-    if (spellId)
-    {
-        switch (spellId)
-        {
-            case 93321: // Control Pet
-                return getClass() == CLASS_HUNTER;
-            case 93375: // Control Demon
-                return getClass() == CLASS_WARLOCK;
-        }
-    }
-    else
-    {
-        if (getClass() == CLASS_HUNTER && !HasAura(93321))
-            return false;
-        if (getClass() == CLASS_WARLOCK && !HasAura(93375))
-            return false;
-    }
-
-    return true;
-}
-
 void Player::PetSpellInitialize()
 {
     // Do not send the pet action bar when Hunters or Warlocks cannot control their pet
diff --git a/src/server/game/Entities/Player/Player.h b/src/server/game/Entities/Player/Player.h
index 3d30b956e1..e0f08f2d95 100644
--- a/src/server/game/Entities/Player/Player.h
+++ b/src/server/game/Entities/Player/Player.h
@@ -1563,13 +1563,12 @@ class TC_GAME_API Player : public Unit, public GridObject<Player>
         bool RemoveMItem(uint32 id);
 
         // Pets
-        void SendPetSpellsMessage(NewPet* pet);
+        void SendPetSpellsMessage(NewPet* pet, bool remove = false);
         void SetCanControlClassPets();
         bool CanControlClassPets() const;
 
     public:
         void SendOnCancelExpectedVehicleRideAura() const;
-        bool CanControlPet(uint32 spellId = 0) const;
         void PetSpellInitialize();
 
         void CharmSpellInitialize();
@@ -2392,12 +2391,6 @@ class TC_GAME_API Player : public Unit, public GridObject<Player>
         // Unsummons the currently active pet and removes the player pet data from its container. The database data will be erased on the next save cycle.
         void AbandonPet();
 
-        PlayerPetData* GetPlayerPetDataCurrent();
-        Optional<uint8> GetFirstUnusedActivePetSlot();
-        Optional<uint8> GetFirstUnusedPetSlot();
-        void DeleteFromPlayerPetDataStore(uint32 petNumber);
-        void AddToPlayerPetDataStore(PlayerPetData&& playerPetData);
-
     protected:
         // Gamemaster whisper whitelist
         GuidList WhisperList;
diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index 40a604d90e..18f25fa00b 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -5583,7 +5583,7 @@ void Unit::SetActivelyControlledSummon(NewPet* pet, bool apply)
     }
 
     if (IsPlayer())
-        ToPlayer()->SendPetSpellsMessage(pet);
+        ToPlayer()->SendPetSpellsMessage(pet, !apply);
 }
 
 Player* Unit::GetControllingPlayer() const
diff --git a/src/server/game/Handlers/PetHandler.cpp b/src/server/game/Handlers/PetHandler.cpp
index 64b88014c3..cddf334bb3 100644
--- a/src/server/game/Handlers/PetHandler.cpp
+++ b/src/server/game/Handlers/PetHandler.cpp
@@ -78,6 +78,7 @@ void HandlePetActionHelper(NewPet* pet, Player* owner, ObjectGuid targetGuid, ui
                     printf("COMMAND_ATTACK -- %u\n", actionValue);
                     break;
                 case COMMAND_ABANDON:
+                    printf("COMMAND_ABANDON -- %u\n", actionValue);
                     if (!pet->HasByteFlag(UNIT_FIELD_BYTES_2, UNIT_BYTES_2_OFFSET_PET_FLAGS, UNIT_CAN_BE_ABANDONED))
                         owner->DismissPet(false);
                     break;
@@ -712,7 +713,7 @@ void WorldSession::HandlePetAbandon(WorldPackets::Pet::PetAbandon& packet)
         return;
 
     NewPet* pet = _player->GetActivelyControlledSummon();
-    if (pet->GetGUID() != packet.Pet)
+    if (!pet || pet->GetGUID() != packet.Pet)
         return;
 
     _player->AbandonPet();
diff --git a/src/server/game/Server/Packets/PetPackets.cpp b/src/server/game/Server/Packets/PetPackets.cpp
index b83550e7b6..f2a3165cf0 100644
--- a/src/server/game/Server/Packets/PetPackets.cpp
+++ b/src/server/game/Server/Packets/PetPackets.cpp
@@ -84,6 +84,9 @@ WorldPacket const* WorldPackets::Pet::PetAdded::Write()
 WorldPacket const* WorldPackets::Pet::PetSpellsMessage::Write()
 {
     _worldPacket << PetGUID;
+    if (PetGUID.IsEmpty())
+        return &_worldPacket;
+
     _worldPacket << uint16(_CreatureFamily);
     _worldPacket << uint32(TimeLimit);
     _worldPacket << uint8(ReactState);
diff --git a/src/server/scripts/Spells/spell_generic.cpp b/src/server/scripts/Spells/spell_generic.cpp
index 7cf63f2f6b..99128829c0 100644
--- a/src/server/scripts/Spells/spell_generic.cpp
+++ b/src/server/scripts/Spells/spell_generic.cpp
@@ -5608,6 +5608,48 @@ class spell_gen_dalaran_shop_keeper_greeting_ping : public SpellScript
     }
 };
 
+enum ControlPet
+{
+    SPELL_CONTROL_DEMON_EFFECT  = 93376,
+    SPELL_CONTROL_PET_EFFECT    = 93322,
+};
+
+// 93321 - Control Pet (Passive)
+// 93375 - Control Demon (Passive)
+class spell_gen_control_pet : public AuraScript
+{
+public:
+    spell_gen_control_pet(uint32 effectSpellId) : AuraScript(), _effectSpellId(effectSpellId) { }
+
+    bool Load() override
+    {
+        return GetCaster()->IsPlayer();
+    }
+
+    bool Validate(SpellInfo const* /*spellInfo*/) override
+    {
+        return ValidateSpellInfo({ _effectSpellId });
+    }
+
+    void AfterApply(AuraEffect const* /*aurEff*/, AuraEffectHandleModes /*mode*/)
+    {
+        Player* player = GetTarget()->ToPlayer();
+        if (!player)
+            return;
+
+        if (!player->CanControlClassPets())
+            player->CastSpell(nullptr, _effectSpellId, TRIGGERED_FULL_MASK);
+    }
+
+    void Register() override
+    {
+        AfterEffectApply.Register(&spell_gen_control_pet::AfterApply, EFFECT_0, SPELL_AURA_DUMMY, AURA_EFFECT_HANDLE_REAL_OR_REAPPLY_MASK);
+    }
+
+private:
+    uint32 _effectSpellId = 0;
+};
+
 void AddSC_generic_spell_scripts()
 {
     new spell_gen_absorb0_hitlimit1();
@@ -5750,4 +5792,6 @@ void AddSC_generic_spell_scripts()
     RegisterSpellScript(spell_gen_shadowmeld);
     RegisterSpellScript(spell_gen_vehicle_control_link);
     RegisterSpellScript(spell_gen_polymorph_cast_visual);
+    RegisterSpellScriptWithArgs(spell_gen_control_pet, "spell_gen_control_demon", SPELL_CONTROL_DEMON_EFFECT);
+    RegisterSpellScriptWithArgs(spell_gen_control_pet, "spell_gen_control_pet", SPELL_CONTROL_PET_EFFECT);
 }
diff --git a/src/server/scripts/Spells/spell_hunter.cpp b/src/server/scripts/Spells/spell_hunter.cpp
index 77a2eeb1d3..0eb468223c 100644
--- a/src/server/scripts/Spells/spell_hunter.cpp
+++ b/src/server/scripts/Spells/spell_hunter.cpp
@@ -746,12 +746,12 @@ class spell_hun_tame_beast : public SpellScript
         if (player->getClass() != CLASS_HUNTER)
             return SendTameFailResult(PET_TAME_FAILURE_CANNOT_TAME_CREATURES);
 
-        if (!player->GetFirstUnusedActivePetSlot())
-            return SendTameFailResult(PET_TAME_FAILURE_TOO_MANY_PETS);
+        //if (!player->GetFirstUnusedActivePetSlot())
+        //    return SendTameFailResult(PET_TAME_FAILURE_TOO_MANY_PETS);
 
-        if (Optional<uint8> slot = player->GetFirstUnusedActivePetSlot())
-            if (!player->HasSpell(callPetSpellIdBySlot[*slot]))
-                return SendTameFailResult(PET_TAME_FAILURE_SLOT_LOCKED);
+        //if (Optional<uint8> slot = player->GetFirstUnusedActivePetSlot())
+         //   if (!player->HasSpell(callPetSpellIdBySlot[*slot]))
+          //      return SendTameFailResult(PET_TAME_FAILURE_SLOT_LOCKED);
 
         if (Creature* target = GetExplTargetUnit()->ToCreature())
         {