Merge pull request #16644 from Treeston/3.3.5-customcharmai

Core/UnitAI: Rework creature-controlled player behavior.
(cherry picked from commit 4e4b2b9a138650ebe6087e18dc392f08dafe4e37)

1 files changed, 30 insertions, 2 deletions

diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index e2b01495e4b..ae0b55352a6 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -44,6 +44,7 @@
 #include "PetAI.h"
 #include "Pet.h"
 #include "Player.h"
+#include "PlayerAI.h"
 #include "QuestDef.h"
 #include "ReputationMgr.h"
 #include "SpellAuraEffects.h"
@@ -14100,11 +14101,21 @@ bool Unit::SetCharmedBy(Unit* charmer, CharmType type, AuraApplication const* au
         ToCreature()->AI()->OnCharmed(true);
         GetMotionMaster()->MoveIdle();
     }
-    else
+    else if (Player* player = ToPlayer())
     {
-        Player* player = ToPlayer();
         if (player->isAFK())
             player->ToggleAFK();
+
+        if (Creature* creatureCharmer = charmer->ToCreature()) // we are charmed by a creature
+        {
+            IsAIEnabled = true;
+            i_disabledAI = i_AI;
+            // set our AI to the charmer's custom charm AI if applicable
+            if (PlayerAI* charmAI = creatureCharmer->AI()->GetAIForCharmedPlayer(player))
+                i_AI = charmAI;
+            else // otherwise use the default charmed player AI
+                i_AI = new SimpleCharmedPlayerAI(player);
+        }
         player->SetClientControl(this, false);
     }
 
@@ -14262,7 +14273,24 @@ void Unit::RemoveCharmedBy(Unit* charmer)
     }
 
     if (Player* player = ToPlayer())
+    {
+        if (charmer->GetTypeId() == TYPEID_UNIT) // charmed by a creature, this means we had PlayerAI
+        {
+            if (i_AI)
+            {
+                // allow charmed player AI to clean up
+                i_AI->OnCharmed(false);
+                // then delete it
+                delete i_AI;
+                // and restore our previous playerAI (if we had one)
+                if ((i_AI = i_disabledAI))
+                    i_disabledAI = nullptr;
+                else
+                    IsAIEnabled = false;
+            }
+        }
         player->SetClientControl(this, true);
+    }
 
     // a guardian should always have charminfo
     if (playerCharmer && this != charmer->GetFirstControlled())