Misc: improve UNIT_STATE_DISTRACTED handling

Author: ccrs

src/server/game/AI/CoreAI/UnitAI.cpp

@@ -30,15 +30,7 @@
 void UnitAI::AttackStart(Unit* victim)
 {
     if (victim && me->Attack(victim, true))
-    {
-        // Clear distracted state on attacking
-        if (me->HasUnitState(UNIT_STATE_DISTRACTED))
-        {
-            me->ClearUnitState(UNIT_STATE_DISTRACTED);
-            me->GetMotionMaster()->Clear();
-        }
         me->GetMotionMaster()->MoveChase(victim);
-    }
 }
 
 void UnitAI::InitializeAI()

src/server/game/Entities/Unit/Unit.cpp

@@ -8773,6 +8773,12 @@ void Unit::AtExitCombat()
     RemoveAurasWithInterruptFlags(AURA_INTERRUPT_FLAG_LEAVE_COMBAT);
 }
 
+void Unit::AtEngage(Unit *)
+{
+    if (HasUnitState(UNIT_STATE_DISTRACTED))
+        GetMotionMaster()->Remove(DISTRACT_MOTION_TYPE);
+}
+
 void Unit::AtTargetAttacked(Unit* target, bool canInitialAggro)
 {
     if (!target->IsEngaged() && !canInitialAggro)

src/server/game/Entities/Unit/Unit.h

@@ -1914,7 +1914,7 @@ class TC_GAME_API Unit : public WorldObject
         virtual void AtEnterCombat() { }
         virtual void AtExitCombat();
 
-        virtual void AtEngage(Unit* /*target*/) {}
+        void AtEngage(Unit* /*target*/);
         virtual void AtDisengage() {}
 
     private:

src/server/game/Movement/MovementGenerators/IdleMovementGenerator.cpp

@@ -140,7 +140,7 @@ MovementGeneratorType RotateMovementGenerator::GetMovementGeneratorType() const
 
 //----------------------------------------------------//
 
-DistractMovementGenerator::DistractMovementGenerator(uint32 timer, float orientation) : _timer(timer), _orientation(orientation)
+DistractMovementGenerator::DistractMovementGenerator(uint32 timer, float orientation) : _timer(timer), _orientation(orientation), _originalOrientation(0.f)
 {
     Mode = MOTION_MODE_DEFAULT;
     Priority = MOTION_PRIORITY_HIGHEST;
@@ -157,6 +157,8 @@ void DistractMovementGenerator::Initialize(Unit* owner)
     if (!owner->IsStandState())
         owner->SetStandState(UNIT_STAND_STATE_STAND);
 
+    _originalOrientation = owner->GetOrientation();
+
     Movement::MoveSplineInit init(owner);
     init.MoveTo(PositionToVector3(*owner), false);
     if (owner->HasUnitMovementFlag(MOVEMENTFLAG_ONTRANSPORT) && !owner->GetTransGUID().IsEmpty())
@@ -196,13 +198,8 @@ void DistractMovementGenerator::Finalize(Unit* owner, bool/* active*/, bool move
 {
     AddFlag(MOVEMENTGENERATOR_FLAG_FINALIZED);
 
-    // TODO: This code should be handled somewhere else
-    // If this is a creature, then return orientation to original position (for idle movement creatures)
     if (movementInform && HasFlag(MOVEMENTGENERATOR_FLAG_INFORM_ENABLED) && owner->GetTypeId() == TYPEID_UNIT)
-    {
-        float angle = owner->ToCreature()->GetHomePosition().GetOrientation();
-        owner->SetFacingTo(angle);
-    }
+        owner->SetFacingTo(_originalOrientation, true);
 }
 
 MovementGeneratorType DistractMovementGenerator::GetMovementGeneratorType() const

src/server/game/Movement/MovementGenerators/IdleMovementGenerator.h

@@ -68,6 +68,7 @@ class DistractMovementGenerator : public MovementGenerator
     private:
         uint32 _timer;
         float _orientation;
+        float _originalOrientation;
 };
 
 class AssistanceDistractMovementGenerator : public DistractMovementGenerator