From ded3bd8900e1b3b422c8baec88f82a38b9308115 Mon Sep 17 00:00:00 2001
From: Nightprince <keramatjokar8855@gmail.com>
Date: Mon, 18 Nov 2019 18:35:25 +0330
Subject: [PATCH] fix(Core/PetAI): aggro and movement of Warlock's IMP (#2379)

---
 src/server/game/AI/CoreAI/PetAI.cpp | 52 ++++++++++++++++++++++++++---
 src/server/game/AI/CoreAI/PetAI.h   | 22 +++++++++++-
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/src/server/game/AI/CoreAI/PetAI.cpp b/src/server/game/AI/CoreAI/PetAI.cpp
index 9ca9cbae2..d6b75c21d 100644
--- a/src/server/game/AI/CoreAI/PetAI.cpp
+++ b/src/server/game/AI/CoreAI/PetAI.cpp
@@ -87,11 +87,52 @@ void PetAI::_doMeleeAttack()
     DoMeleeAttackIfReady();
 }
 
-bool PetAI::_canMeleeAttack() const
+bool PetAI::_canMeleeAttack()
 {
-    return me->GetEntry() != 416 /*ENTRY_IMP*/ &&
-        me->GetEntry() != 510 /*ENTRY_WATER_ELEMENTAL*/ &&
-        me->GetEntry() != 37994 /*ENTRY_WATER_ELEMENTAL_PERM*/; 
+    combatRange = 0.f;
+    switch (me->GetEntry())
+    {
+        case ENTRY_IMP:
+        case ENTRY_WATER_ELEMENTAL:
+        case ENTRY_WATER_ELEMENTAL_PERM:
+        {
+            for (uint8 i = 0; i < me->GetPetAutoSpellSize(); ++i)
+            {
+                uint32 spellID = me->GetPetAutoSpellOnPos(i);
+                switch (spellID)
+                {
+                    case IMP_FIREBOLT_RANK_1:
+                    case IMP_FIREBOLT_RANK_2:
+                    case IMP_FIREBOLT_RANK_3:
+                    case IMP_FIREBOLT_RANK_4:
+                    case IMP_FIREBOLT_RANK_5:
+                    case IMP_FIREBOLT_RANK_6:
+                    case IMP_FIREBOLT_RANK_7:
+                    case IMP_FIREBOLT_RANK_8:
+                    case IMP_FIREBOLT_RANK_9:
+                    case WATER_ELEMENTAL_WATERBOLT_1:
+                    case WATER_ELEMENTAL_WATERBOLT_2:
+                    {
+                        SpellInfo const* spellInfo = sSpellMgr->GetSpellInfo(spellID);
+                        int32 mana = me->GetPower(POWER_MANA);
+
+                        if (mana >= spellInfo->CalcPowerCost(me, spellInfo->GetSchoolMask()))
+                        {
+                            combatRange = spellInfo->GetMaxRange();
+                            return true;
+                        }
+                    }
+                    default:
+                        break;
+                }
+            }
+            return false;
+        }
+        default:
+            break;
+    }
+
+    return true;
 }
 
 void PetAI::UpdateAI(uint32 diff)
@@ -563,7 +604,8 @@ void PetAI::DoAttack(Unit* target, bool chase)
 
         if (chase)
         {
-            me->GetMotionMaster()->MoveChase(target, !_canMeleeAttack() ? 20.0f: 0.0f, me->GetAngle(target));
+            if (_canMeleeAttack())
+                me->GetMotionMaster()->MoveChase(target, combatRange, float(M_PI));
         }
         else // (Stay && ((Aggressive || Defensive) && In Melee Range)))
         {
diff --git a/src/server/game/AI/CoreAI/PetAI.h b/src/server/game/AI/CoreAI/PetAI.h
index 0dca22d5b..c4f42ba63 100644
--- a/src/server/game/AI/CoreAI/PetAI.h
+++ b/src/server/game/AI/CoreAI/PetAI.h
@@ -13,6 +13,25 @@
 class Creature;
 class Spell;
 
+enum SpecialPets
+{
+    ENTRY_IMP                   =   416,
+    ENTRY_WATER_ELEMENTAL       =   510,
+    ENTRY_WATER_ELEMENTAL_PERM  = 37994,
+
+    IMP_FIREBOLT_RANK_1         =  3110,
+    IMP_FIREBOLT_RANK_2         =  7799,
+    IMP_FIREBOLT_RANK_3         =  7800,
+    IMP_FIREBOLT_RANK_4         =  7801,
+    IMP_FIREBOLT_RANK_5         =  7802,
+    IMP_FIREBOLT_RANK_6         = 11762,
+    IMP_FIREBOLT_RANK_7         = 11763,
+    IMP_FIREBOLT_RANK_8         = 27267,
+    IMP_FIREBOLT_RANK_9         = 47964,
+    WATER_ELEMENTAL_WATERBOLT_1 = 31707,
+    WATER_ELEMENTAL_WATERBOLT_2 = 72898
+};
+
 class PetAI : public CreatureAI
 {
     public:
@@ -43,13 +62,14 @@ class PetAI : public CreatureAI
         bool _needToStop(void);
         void _stopAttack(void);
         void _doMeleeAttack();
-        bool _canMeleeAttack() const;
+        bool _canMeleeAttack();
 
         void UpdateAllies();
 
         TimeTracker i_tracker;
         std::set<uint64> m_AllySet;
         uint32 m_updateAlliesTimer;
+        float combatRange;
 
         Unit* SelectNextTarget(bool allowAutoSelect) const;
         void HandleReturnMovement();