Jetpack QOL tweaks (#9563)

Content.Shared/Movement/Systems/SharedJetpackSystem.cs

@@ -4,17 +4,24 @@ using Content.Shared.Interaction.Events;
 using Content.Shared.Inventory;
 using Content.Shared.Movement.Components;
 using Content.Shared.Movement.Events;
+using Content.Shared.Popups;
 using Robust.Shared.Containers;
+using Robust.Shared.GameStates;
 using Robust.Shared.Map;
+using Robust.Shared.Network;
+using Robust.Shared.Player;
 using Robust.Shared.Serialization;
+using Robust.Shared.Timing;
 
 namespace Content.Shared.Movement.Systems;
 
 public abstract class SharedJetpackSystem : EntitySystem
 {
+    [Dependency] private readonly IGameTiming _timing = default!;
+    [Dependency] private readonly INetManager _network = default!;
     [Dependency] protected readonly SharedContainerSystem Container = default!;
-    [Dependency] protected readonly IMapManager MapManager = default!;
-    [Dependency] protected readonly MovementSpeedModifierSystem MovementSpeedModifier = default!;
+    [Dependency] private readonly MovementSpeedModifierSystem _movementSpeedModifier = default!;
+    [Dependency] private readonly SharedPopupSystem _popups = default!;
 
     public override void Initialize()
     {
@@ -25,9 +32,44 @@ public abstract class SharedJetpackSystem : EntitySystem
         SubscribeLocalEvent<JetpackUserComponent, CanWeightlessMoveEvent>(OnJetpackUserCanWeightless);
         SubscribeLocalEvent<JetpackUserComponent, MobMovementProfileEvent>(OnJetpackUserMovement);
         SubscribeLocalEvent<JetpackUserComponent, EntParentChangedMessage>(OnJetpackUserEntParentChanged);
+        SubscribeLocalEvent<JetpackUserComponent, ComponentGetState>(OnJetpackUserGetState);
+        SubscribeLocalEvent<JetpackUserComponent, ComponentHandleState>(OnJetpackUserHandleState);
+
         SubscribeLocalEvent<GravityChangedMessage>(OnJetpackUserGravityChanged);
     }
 
+    private void OnJetpackUserGravityChanged(GravityChangedMessage ev)
+    {
+        var gridUid = ev.ChangedGridIndex;
+        var jetpackQuery = GetEntityQuery<JetpackComponent>();
+
+        foreach (var (user, transform) in EntityQuery<JetpackUserComponent, TransformComponent>(true))
+        {
+            if (transform.GridUid == gridUid && ev.HasGravity &&
+                jetpackQuery.TryGetComponent(user.Jetpack, out var jetpack))
+            {
+                if (_timing.IsFirstTimePredicted)
+                    _popups.PopupEntity(Loc.GetString("jetpack-to-grid"), user.Jetpack, Filter.Entities(user.Owner));
+
+                SetEnabled(jetpack, false, user.Owner);
+            }
+        }
+    }
+
+    private void OnJetpackUserHandleState(EntityUid uid, JetpackUserComponent component, ref ComponentHandleState args)
+    {
+        if (args.Current is not JetpackUserComponentState state) return;
+        component.Jetpack = state.Jetpack;
+    }
+
+    private void OnJetpackUserGetState(EntityUid uid, JetpackUserComponent component, ref ComponentGetState args)
+    {
+        args.State = new JetpackUserComponentState()
+        {
+            Jetpack = component.Jetpack,
+        };
+    }
+
     private void OnJetpackDropped(EntityUid uid, JetpackComponent component, DroppedEvent args)
     {
         SetEnabled(component, false, args.User);
@@ -51,16 +93,13 @@ public abstract class SharedJetpackSystem : EntitySystem
 
     private void OnJetpackUserEntParentChanged(EntityUid uid, JetpackUserComponent component, ref EntParentChangedMessage args)
     {
-        MovementSpeedModifier.RefreshMovementSpeedModifiers(component.Owner);
-    }
-
-    private void OnJetpackUserGravityChanged(GravityChangedMessage ev)
-    {
-        var gridUid = ev.ChangedGridIndex;
-        foreach (var (_, transform) in EntityQuery<JetpackUserComponent, TransformComponent>(true))
+        if (TryComp<JetpackComponent>(component.Jetpack, out var jetpack) &&
+            !CanEnableOnGrid(args.Transform.GridUid))
         {
-            if(transform.GridUid == gridUid)
-                MovementSpeedModifier.RefreshMovementSpeedModifiers(transform.Owner);
+            SetEnabled(jetpack, false, uid);
+
+            if (_timing.IsFirstTimePredicted && _network.IsClient)
+                _popups.PopupEntity(Loc.GetString("jetpack-to-grid"), uid, Filter.Entities(uid));
         }
     }
 
@@ -70,15 +109,31 @@ public abstract class SharedJetpackSystem : EntitySystem
         user.Acceleration = component.Acceleration;
         user.Friction = component.Friction;
         user.WeightlessModifier = component.WeightlessModifier;
+        user.Jetpack = component.Owner;
     }
 
     private void OnJetpackToggle(EntityUid uid, JetpackComponent component, ToggleJetpackEvent args)
     {
         if (args.Handled) return;
 
+        if (TryComp<TransformComponent>(uid, out var xform) && !CanEnableOnGrid(xform.GridUid))
+        {
+            if (_timing.IsFirstTimePredicted)
+                _popups.PopupEntity(Loc.GetString("jetpack-no-station"), uid, Filter.Entities(args.Performer));
+
+            return;
+        }
+
         SetEnabled(component, !IsEnabled(uid));
     }
 
+    private bool CanEnableOnGrid(EntityUid? gridUid)
+    {
+        return gridUid == null ||
+               (TryComp<GravityComponent>(gridUid, out var gravity) &&
+                !gravity.Enabled);
+    }
+
     private void OnJetpackGetAction(EntityUid uid, JetpackComponent component, GetItemActionsEvent args)
     {
         args.Actions.Add(component.ToggleAction);
@@ -123,7 +178,7 @@ public abstract class SharedJetpackSystem : EntitySystem
                 RemComp<JetpackUserComponent>(user.Value);
             }
 
-            MovementSpeedModifier.RefreshMovementSpeedModifiers(user.Value);
+            _movementSpeedModifier.RefreshMovementSpeedModifiers(user.Value);
         }
 
         TryComp<AppearanceComponent>(component.Owner, out var appearance);
@@ -133,17 +188,18 @@ public abstract class SharedJetpackSystem : EntitySystem
 
     public bool IsUserFlying(EntityUid uid)
     {
-        return HasComp<JetpackUserComponent>(uid) &&
-               TryComp<PhysicsComponent>(uid, out var physicsComponent) &&
-               uid.IsWeightless(physicsComponent, mapManager: MapManager, entityManager: EntityManager);
+        return HasComp<JetpackUserComponent>(uid);
     }
 
-    protected abstract bool CanEnable(JetpackComponent component);
+    protected virtual bool CanEnable(JetpackComponent component)
+    {
+        return true;
+    }
 
     [Serializable, NetSerializable]
-    protected sealed class JetpackComponentState : ComponentState
+    protected sealed class JetpackUserComponentState : ComponentState
     {
-        public bool Enabled;
+        public EntityUid Jetpack;
     }
 }