add StationTrackerComponent (#36803)

* maybe I am cooking

* logmissing

* copy paste oops

* add some stuff

* review

* fix

* rerun tests

---------

Co-authored-by: slarticodefast <161409025+slarticodefast@users.noreply.github.com>

Content.Client/Station/StationSystem.cs

@@ -5,7 +5,7 @@ namespace Content.Client.Station;
 /// <summary>
 /// This handles letting the client know stations are a thing. Only really used by an admin menu.
 /// </summary>
-public sealed class StationSystem : EntitySystem
+public sealed partial class StationSystem : SharedStationSystem
 {
     private readonly List<(string Name, NetEntity Entity)> _stations = new();
 
@@ -15,11 +15,14 @@ public sealed class StationSystem : EntitySystem
     /// <remarks>
     /// I'd have this just invoke an entity query, but we're on the client and the client barely knows about stations.
     /// </remarks>
+    // TODO: Stations have a global PVS override now, this can probably be changed into a query.
     public IReadOnlyList<(string Name, NetEntity Entity)> Stations => _stations;
 
     /// <inheritdoc/>
     public override void Initialize()
     {
+        base.Initialize();
+
         SubscribeNetworkEvent<StationsUpdatedEvent>(StationsUpdated);
     }
 

Content.Server/Station/Systems/StationSystem.cs

@@ -27,7 +27,7 @@ namespace Content.Server.Station.Systems;
 /// For jobs, look at StationJobSystem. For spawning, look at StationSpawningSystem.
 /// </summary>
 [PublicAPI]
-public sealed class StationSystem : EntitySystem
+public sealed partial class StationSystem : SharedStationSystem
 {
     [Dependency] private readonly ILogManager _logManager = default!;
     [Dependency] private readonly IPlayerManager _player = default!;
@@ -48,6 +48,8 @@ public sealed class StationSystem : EntitySystem
     /// <inheritdoc/>
     public override void Initialize()
     {
+        base.Initialize();
+
         _sawmill = _logManager.GetSawmill("station");
 
         _gridQuery = GetEntityQuery<MapGridComponent>();
@@ -60,6 +62,9 @@ public sealed class StationSystem : EntitySystem
         SubscribeLocalEvent<StationMemberComponent, ComponentShutdown>(OnStationGridDeleted);
         SubscribeLocalEvent<StationMemberComponent, PostGridSplitEvent>(OnStationSplitEvent);
 
+        SubscribeLocalEvent<StationGridAddedEvent>(OnStationGridAdded);
+        SubscribeLocalEvent<StationGridRemovedEvent>(OnStationGridRemoved);
+
         _player.PlayerStatusChanged += OnPlayerStatusChanged;
     }
 
@@ -90,6 +95,18 @@ public sealed class StationSystem : EntitySystem
         }
     }
 
+    private void UpdateTrackersOnGrid(EntityUid gridId, EntityUid? station)
+    {
+        var query = EntityQueryEnumerator<StationTrackerComponent, TransformComponent>();
+        while (query.MoveNext(out var uid, out var tracker, out var xform))
+        {
+            if (xform.GridUid == gridId)
+            {
+                SetStation((uid, tracker), station);
+            }
+        }
+    }
+
     #region Event handlers
 
     private void OnStationAdd(EntityUid uid, StationDataComponent component, ComponentStartup args)
@@ -107,6 +124,9 @@ public sealed class StationSystem : EntitySystem
         foreach (var grid in component.Grids)
         {
             RemComp<StationMemberComponent>(grid);
+
+            // If the station gets deleted, we raise the event for every grid that was a part of it
+            RaiseLocalEvent(new StationGridRemovedEvent(grid, uid));
         }
 
         RaiseNetworkEvent(new StationsUpdatedEvent(GetStationNames()), Filter.Broadcast());
@@ -159,6 +179,18 @@ public sealed class StationSystem : EntitySystem
         }
     }
 
+    private void OnStationGridAdded(StationGridAddedEvent ev)
+    {
+        // When a grid is added to a station, update all trackers on that grid
+        UpdateTrackersOnGrid(ev.GridId, ev.Station);
+    }
+
+    private void OnStationGridRemoved(StationGridRemovedEvent ev)
+    {
+        // When a grid is removed from a station, update all trackers on that grid to null
+        UpdateTrackersOnGrid(ev.GridId, null);
+    }
+
     #endregion Event handlers
 
     /// <summary>
@@ -353,7 +385,7 @@ public sealed class StationSystem : EntitySystem
         stationMember.Station = station;
         stationData.Grids.Add(mapGrid);
 
-        RaiseLocalEvent(station, new StationGridAddedEvent(mapGrid, false), true);
+        RaiseLocalEvent(station, new StationGridAddedEvent(mapGrid, station, false), true);
 
         _sawmill.Info($"Adding grid {mapGrid} to station {Name(station)} ({station})");
     }
@@ -376,7 +408,7 @@ public sealed class StationSystem : EntitySystem
         RemComp<StationMemberComponent>(mapGrid);
         stationData.Grids.Remove(mapGrid);
 
-        RaiseLocalEvent(station, new StationGridRemovedEvent(mapGrid), true);
+        RaiseLocalEvent(station, new StationGridRemovedEvent(mapGrid, station), true);
         _sawmill.Info($"Removing grid {mapGrid} from station {Name(station)} ({station})");
     }
 
@@ -553,15 +585,21 @@ public sealed class StationGridAddedEvent : EntityEventArgs
     /// </summary>
     public EntityUid GridId;
 
+    /// <summary>
+    /// EntityUid of the station this grid was added to.
+    /// </summary>
+    public EntityUid Station;
+
     /// <summary>
     /// Indicates that the event was fired during station setup,
     /// so that it can be ignored if StationInitializedEvent was already handled.
     /// </summary>
     public bool IsSetup;
 
-    public StationGridAddedEvent(EntityUid gridId, bool isSetup)
+    public StationGridAddedEvent(EntityUid gridId, EntityUid station, bool isSetup)
     {
         GridId = gridId;
+        Station = station;
         IsSetup = isSetup;
     }
 }
@@ -577,9 +615,15 @@ public sealed class StationGridRemovedEvent : EntityEventArgs
     /// </summary>
     public EntityUid GridId;
 
-    public StationGridRemovedEvent(EntityUid gridId)
+    /// <summary>
+    /// EntityUid of the station this grid was added to.
+    /// </summary>
+    public EntityUid Station;
+
+    public StationGridRemovedEvent(EntityUid gridId, EntityUid station)
     {
         GridId = gridId;
+        Station = station;
     }
 }
 

Content.Shared/Station/Components/StationTrackerComponent.cs

@@ -0,0 +1,18 @@
+using Robust.Shared.GameStates;
+
+namespace Content.Shared.Station.Components;
+
+/// <summary>
+/// Component that tracks which station an entity is currently on.
+/// Mainly used for UI purposes on the client to easily get station-specific data like alert levels.
+/// </summary>
+[RegisterComponent, NetworkedComponent, AutoGenerateComponentState, Access(typeof(SharedStationSystem))]
+public sealed partial class StationTrackerComponent : Component
+{
+    /// <summary>
+    /// The station this entity is currently on, if any.
+    /// Null when in space or not on any grid.
+    /// </summary>
+    [DataField, AutoNetworkedField]
+    public EntityUid? Station;
+}

Content.Shared/Station/SharedStationSystem.cs

@@ -0,0 +1,111 @@
+using Content.Shared.Station.Components;
+using JetBrains.Annotations;
+using Robust.Shared.Map;
+
+namespace Content.Shared.Station;
+
+public abstract partial class SharedStationSystem : EntitySystem
+{
+    [Dependency] private readonly MetaDataSystem _meta = default!;
+
+    private EntityQuery<TransformComponent> _xformQuery;
+    private EntityQuery<StationMemberComponent> _stationMemberQuery;
+
+    /// <inheritdoc/>
+    public override void Initialize()
+    {
+        base.Initialize();
+
+        _xformQuery = GetEntityQuery<TransformComponent>();
+        _stationMemberQuery = GetEntityQuery<StationMemberComponent>();
+
+        SubscribeLocalEvent<StationTrackerComponent, MapInitEvent>(OnTrackerMapInit);
+        SubscribeLocalEvent<StationTrackerComponent, ComponentRemove>(OnTrackerRemove);
+        SubscribeLocalEvent<StationTrackerComponent, GridUidChangedEvent>(OnTrackerGridChanged);
+        SubscribeLocalEvent<StationTrackerComponent, MetaFlagRemoveAttemptEvent>(OnMetaFlagRemoveAttempt);
+    }
+
+    private void OnTrackerMapInit(Entity<StationTrackerComponent> ent, ref MapInitEvent args)
+    {
+        _meta.AddFlag(ent, MetaDataFlags.ExtraTransformEvents);
+        UpdateStationTracker(ent.AsNullable());
+    }
+
+    private void OnTrackerRemove(Entity<StationTrackerComponent> ent, ref ComponentRemove args)
+    {
+        _meta.RemoveFlag(ent, MetaDataFlags.ExtraTransformEvents);
+    }
+
+    private void OnTrackerGridChanged(Entity<StationTrackerComponent> ent, ref GridUidChangedEvent args)
+    {
+        UpdateStationTracker((ent, ent.Comp, args.Transform));
+    }
+
+    private void OnMetaFlagRemoveAttempt(Entity<StationTrackerComponent> ent, ref MetaFlagRemoveAttemptEvent args)
+    {
+        if ((args.ToRemove & MetaDataFlags.ExtraTransformEvents) != 0 &&
+            ent.Comp.LifeStage <= ComponentLifeStage.Running)
+        {
+            args.ToRemove &= ~MetaDataFlags.ExtraTransformEvents;
+        }
+    }
+
+    /// <summary>
+    /// Updates the station tracker component based on entity's current location.
+    /// </summary>
+    [PublicAPI]
+    public void UpdateStationTracker(Entity<StationTrackerComponent?, TransformComponent?> ent)
+    {
+        if (!Resolve(ent, ref ent.Comp1))
+            return;
+
+        var xform = ent.Comp2;
+
+        if (!_xformQuery.Resolve(ent, ref xform))
+            return;
+
+        // Entity is in nullspace or not on a grid
+        if (xform.MapID == MapId.Nullspace || xform.GridUid == null)
+        {
+            SetStation(ent, null);
+            return;
+        }
+
+        // Check if the grid is part of a station
+        if (!_stationMemberQuery.TryGetComponent(xform.GridUid.Value, out var stationMember))
+        {
+            SetStation(ent, null);
+            return;
+        }
+
+        SetStation(ent, stationMember.Station);
+    }
+
+    /// <summary>
+    /// Sets the station for a StationTrackerComponent.
+    /// </summary>
+    [PublicAPI]
+    public void SetStation(Entity<StationTrackerComponent?> ent, EntityUid? station)
+    {
+        if (!Resolve(ent, ref ent.Comp))
+            return;
+
+        if (ent.Comp.Station == station)
+            return;
+
+        ent.Comp.Station = station;
+        Dirty(ent);
+    }
+
+    /// <summary>
+    /// Gets the station an entity is currently on, if any.
+    /// </summary>
+    [PublicAPI]
+    public EntityUid? GetCurrentStation(Entity<StationTrackerComponent?> ent)
+    {
+        if (!Resolve(ent, ref ent.Comp, logMissing: false))
+            return null;
+
+        return ent.Comp.Station;
+    }
+}