End Levitate when no rounds remaining

Levitate is a ConstantEffect where underlying flag is not cleared by EntityEffectManager in ClearConstantEffects(). This results in Levitate effect persisting when paired with longer-running effects. Call StopLevitating() when there are no RoundsRemaining. Also added a reference cache for LevitateMotor and EnemyMotor to avoid using GetComponent() every call.
Interkarma · Dec 12, 2022 · e18422e · e18422e
1 parent 5c6f55c
commit e18422e
Showing 1 changed file with 39 additions and 9 deletions.
diff --git a/Assets/Scripts/Game/MagicAndEffects/Effects/Thaumaturgy/Levitate.cs b/Assets/Scripts/Game/MagicAndEffects/Effects/Thaumaturgy/Levitate.cs
@@ -22,6 +22,9 @@ public class Levitate : IncumbentEffect
     {
         public static readonly string EffectKey = "Levitate";
 
+        LevitateMotor levitateMotor;
+        EnemyMotor enemyMotor;
+
         public override void SetProperties()
         {
             properties.Key = EffectKey;
@@ -54,7 +57,10 @@ public override void SetPotionProperties()
         public override void ConstantEffect()
         {
             base.ConstantEffect();
-            StartLevitating();
+            if (RoundsRemaining > 0)
+                StartLevitating();
+            else
+                StopLevitating();
         }
 
         public override void Start(EntityEffectManager manager, DaggerfallEntityBehaviour caster = null)
@@ -96,13 +102,11 @@ void StartLevitating()
             // Enable levitation for player or enemies
             if (entityBehaviour.EntityType == EntityTypes.Player)
             {
-                GameManager.Instance.PlayerMotor.GetComponent<LevitateMotor>().IsLevitating = true;
+                SetLevitateMotor(true);
             }
             else if (entityBehaviour.EntityType == EntityTypes.EnemyMonster || entityBehaviour.EntityType == EntityTypes.EnemyClass)
             {
-                EnemyMotor enemyMotor = entityBehaviour.GetComponent<EnemyMotor>();
-                if (enemyMotor)
-                    enemyMotor.IsLevitating = true;
+                SetEnemyMotor(true);
             }
         }
 
@@ -116,13 +120,39 @@ void StopLevitating()
             // Disable levitation for player or enemies
             if (entityBehaviour.EntityType == EntityTypes.Player)
             {
-                GameManager.Instance.PlayerMotor.GetComponent<LevitateMotor>().IsLevitating = GameManager.Instance.PlayerEntity.NoClipMode;
+                SetLevitateMotor(GameManager.Instance.PlayerEntity.NoClipMode);
             }
             else if (entityBehaviour.EntityType == EntityTypes.EnemyMonster || entityBehaviour.EntityType == EntityTypes.EnemyClass)
             {
-                EnemyMotor enemyMotor = entityBehaviour.GetComponent<EnemyMotor>();
-                if (enemyMotor)
-                    enemyMotor.IsLevitating = false;
+                SetEnemyMotor(false);
+            }
+        }
+
+        void SetLevitateMotor(bool state)
+        {
+            if (levitateMotor)
+                levitateMotor.IsLevitating = state;
+            else
+            {
+                levitateMotor = GameManager.Instance.PlayerMotor.GetComponent<LevitateMotor>();
+                if (levitateMotor)
+                    levitateMotor.IsLevitating = state;
+            }
+        }
+
+        void SetEnemyMotor(bool state)
+        {
+            if (enemyMotor)
+                enemyMotor.IsLevitating = state;
+            else
+            {
+                DaggerfallEntityBehaviour entityBehaviour = GetPeeredEntityBehaviour(manager);
+                if (entityBehaviour)
+                {
+                    enemyMotor = entityBehaviour.GetComponent<EnemyMotor>();
+                    if (enemyMotor)
+                        enemyMotor.IsLevitating = state;
+                }
             }
         }
     }