Core/BattlePets: Improvements when enabling battle pets (#27367)

5 files changed, 39 insertions, 17 deletions

diff --git a/src/server/game/BattlePets/BattlePetMgr.cpp b/src/server/game/BattlePets/BattlePetMgr.cpp
index 79754f190aa..0b6f2ebeef7 100644
--- a/src/server/game/BattlePets/BattlePetMgr.cpp
+++ b/src/server/game/BattlePets/BattlePetMgr.cpp
@@ -27,6 +27,7 @@
 #include "ObjectMgr.h"
 #include "Player.h"
 #include "Realm.h"
+#include "Util.h"
 #include "World.h"
 #include "WorldSession.h"
 
@@ -201,7 +202,7 @@ uint32 BattlePetMgr::SelectPetDisplay(BattlePetSpeciesEntry const* speciesEntry)
 BattlePetMgr::BattlePetMgr(WorldSession* owner)
 {
     _owner = owner;
-    for (uint8 i = 0; i < MAX_PET_BATTLE_SLOTS; ++i)
+    for (uint8 i = 0; i < AsUnderlyingType(BattlePetSlot::Count); ++i)
     {
         WorldPackets::BattlePet::BattlePetSlot slot;
         slot.Index = i;
@@ -558,15 +559,19 @@ uint32 BattlePetMgr::GetPetUniqueSpeciesCount() const
     return speciesIds.size();
 }
 
-void BattlePetMgr::UnlockSlot(uint8 slot)
+void BattlePetMgr::UnlockSlot(BattlePetSlot slot)
 {
-    if (!_slots[slot].Locked)
+    if (slot >= BattlePetSlot::Count)
         return;
 
-    _slots[slot].Locked = false;
+    uint8 slotIndex = AsUnderlyingType(slot);
+    if (!_slots[slotIndex].Locked)
+        return;
+
+    _slots[slotIndex].Locked = false;
 
     WorldPackets::BattlePet::PetBattleSlotUpdates updates;
-    updates.Slots.push_back(_slots[slot]);
+    updates.Slots.push_back(_slots[slotIndex]);
     updates.AutoSlotted = false; // what's this?
     updates.NewSlot = true; // causes the "new slot unlocked" bubble to appear
     _owner->SendPacket(updates.Write());
@@ -654,11 +659,11 @@ void BattlePetMgr::SummonPet(ObjectGuid guid)
     if (!speciesEntry)
         return;
 
-    // TODO: set proper CreatureID for spell DEFAULT_SUMMON_BATTLE_PET_SPELL (default EffectMiscValueA is 40721 - Murkimus the Gladiator)
+    // TODO: set proper CreatureID for spell SPELL_SUMMON_BATTLE_PET (default EffectMiscValueA is 40721 - Murkimus the Gladiator)
     Player* player = _owner->GetPlayer();
     player->SetSummonedBattlePetGUID(guid);
     player->SetCurrentBattlePetBreedQuality(pet->PacketInfo.Quality);
-    player->CastSpell(player, speciesEntry->SummonSpellID ? speciesEntry->SummonSpellID : uint32(DEFAULT_SUMMON_BATTLE_PET_SPELL));
+    player->CastSpell(player, speciesEntry->SummonSpellID ? speciesEntry->SummonSpellID : uint32(SPELL_SUMMON_BATTLE_PET));
 }
 
 void BattlePetMgr::DismissPet()
diff --git a/src/server/game/BattlePets/BattlePetMgr.h b/src/server/game/BattlePets/BattlePetMgr.h
index 290a070cdbd..ddfd04f4fec 100644
--- a/src/server/game/BattlePets/BattlePetMgr.h
+++ b/src/server/game/BattlePets/BattlePetMgr.h
@@ -27,11 +27,13 @@ struct BattlePetSpeciesEntry;
 
 enum BattlePetMisc
 {
-    MAX_PET_BATTLE_SLOTS                = 3,
     DEFAULT_MAX_BATTLE_PETS_PER_SPECIES = 3,
     BATTLE_PET_CAGE_ITEM_ID             = 82800,
-    DEFAULT_SUMMON_BATTLE_PET_SPELL     = 118301,
-    SPELL_VISUAL_UNCAGE_PET             = 222
+    SPELL_VISUAL_UNCAGE_PET             = 222,
+
+    SPELL_BATTLE_PET_TRAINING           = 125610,
+    SPELL_REVIVE_BATTLE_PETS            = 125439,
+    SPELL_SUMMON_BATTLE_PET             = 118301
 };
 
 enum class BattlePetBreedQuality : uint8
@@ -69,6 +71,15 @@ enum class BattlePetError : uint8
     TooHighLevelToUncage   = 7  // This pet is too high level for you to uncage.
 };
 
+enum class BattlePetSlot : uint8
+{
+    Slot0 = 0,
+    Slot1 = 1,
+    Slot2 = 2,
+
+    Count
+};
+
 // 6.2.4
 enum FlagsControlType
 {
@@ -145,8 +156,8 @@ public:
     bool HasMaxPetCount(BattlePetSpeciesEntry const* battlePetSpecies, ObjectGuid ownerGuid) const;
     uint32 GetPetUniqueSpeciesCount() const;
 
-    WorldPackets::BattlePet::BattlePetSlot* GetSlot(uint8 slot) { return slot < _slots.size() ? &_slots[slot] : nullptr; }
-    void UnlockSlot(uint8 slot);
+    WorldPackets::BattlePet::BattlePetSlot* GetSlot(BattlePetSlot slot) { return slot < BattlePetSlot::Count ? &_slots[size_t(slot)] : nullptr; }
+    void UnlockSlot(BattlePetSlot slot);
 
     WorldSession* GetOwner() const { return _owner; }
 
@@ -169,6 +180,8 @@ public:
     bool HasJournalLock() const { return _hasJournalLock; }
     void ToggleJournalLock(bool lock) { _hasJournalLock = lock; }
 
+    bool IsBattlePetSystemEnabled() { return GetSlot(BattlePetSlot::Slot0)->Locked != true; }
+
 private:
     WorldSession* _owner;
     bool _hasJournalLock = false;
diff --git a/src/server/game/Entities/Player/Player.cpp b/src/server/game/Entities/Player/Player.cpp
index d0a48fdee6e..efe991c3d0c 100644
--- a/src/server/game/Entities/Player/Player.cpp
+++ b/src/server/game/Entities/Player/Player.cpp
@@ -18616,6 +18616,10 @@ bool Player::LoadFromDB(ObjectGuid guid, CharacterDatabaseQueryHolder* holder)
         _restMgr->AddRestBonus(REST_TYPE_XP, time_diff * _restMgr->CalcExtraPerSec(REST_TYPE_XP, bubble));
     }
 
+    // Unlock battle pet system if it's enabled in bnet account
+    if (GetSession()->GetBattlePetMgr()->IsBattlePetSystemEnabled())
+        LearnSpell(SPELL_BATTLE_PET_TRAINING, false);
+
     m_achievementMgr->CheckAllAchievementCriteria(this);
     m_questObjectiveCriteriaMgr->CheckAllQuestObjectiveCriteria(this);
 
diff --git a/src/server/game/Handlers/BattlePetHandler.cpp b/src/server/game/Handlers/BattlePetHandler.cpp
index 93bca65d3f5..bfe2b2ea36a 100644
--- a/src/server/game/Handlers/BattlePetHandler.cpp
+++ b/src/server/game/Handlers/BattlePetHandler.cpp
@@ -38,7 +38,7 @@ void WorldSession::HandleBattlePetRequestJournalLock(WorldPackets::BattlePet::Ba
 void WorldSession::HandleBattlePetSetBattleSlot(WorldPackets::BattlePet::BattlePetSetBattleSlot& battlePetSetBattleSlot)
 {
     if (BattlePetMgr::BattlePet* pet = GetBattlePetMgr()->GetPet(battlePetSetBattleSlot.PetGuid))
-        if (WorldPackets::BattlePet::BattlePetSlot* slot = GetBattlePetMgr()->GetSlot(battlePetSetBattleSlot.Slot))
+        if (WorldPackets::BattlePet::BattlePetSlot* slot = GetBattlePetMgr()->GetSlot(BattlePetSlot(battlePetSetBattleSlot.Slot)))
             slot->Pet = pet->PacketInfo;
 }
 
diff --git a/src/server/game/Spells/SpellEffects.cpp b/src/server/game/Spells/SpellEffects.cpp
index cc1ea5cc0ed..ffdc611796f 100644
--- a/src/server/game/Spells/SpellEffects.cpp
+++ b/src/server/game/Spells/SpellEffects.cpp
@@ -5498,12 +5498,12 @@ void Spell::EffectEnableBattlePets()
     if (effectHandleMode != SPELL_EFFECT_HANDLE_HIT_TARGET)
         return;
 
-    if (!unitTarget || unitTarget->GetTypeId() != TYPEID_PLAYER)
+    if (!unitTarget || !unitTarget->IsPlayer())
         return;
 
-    Player* plr = unitTarget->ToPlayer();
-    plr->AddPlayerFlag(PLAYER_FLAGS_PET_BATTLES_UNLOCKED);
-    plr->GetSession()->GetBattlePetMgr()->UnlockSlot(0);
+    Player* player = unitTarget->ToPlayer();
+    player->AddPlayerFlag(PLAYER_FLAGS_PET_BATTLES_UNLOCKED);
+    player->GetSession()->GetBattlePetMgr()->UnlockSlot(BattlePetSlot::Slot0);
 }
 
 void Spell::EffectLaunchQuestChoice()