CatStation/Content.Server/Station/Systems/ServerStationSystem.cs

using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using Content.Server.Atmos.EntitySystems;
using Content.Server.Chat.Systems;
using Content.Server.GameTicking;
using Content.Server.Station.Events;
using Content.Shared.GameTicking;
using Content.Shared.Station;
using Content.Shared.Station.Components;
using JetBrains.Annotations;
using Robust.Server.GameStates;
using Robust.Server.Player;
using Robust.Shared.Enums;
using Robust.Shared.Map;
using Robust.Shared.Map.Components;
using Robust.Shared.Player;
using Robust.Shared.Utility;

namespace Content.Server.Station.Systems;

/// <summary>
/// System that manages stations.
/// A station is, by default, just a name, optional map prototype, and optional grids.
/// For jobs, look at StationJobSystem. For spawning, look at StationSpawningSystem.
/// </summary>
[PublicAPI]
public sealed partial class ServerStationSystem : Shared.Station.Systems.StationSystem
{
    [Dependency] private IPlayerManager _player = default!;
    [Dependency] private AtmosphereSystem _atmos = default!;
    [Dependency] private ChatSystem _chatSystem = default!;
    [Dependency] private PvsOverrideSystem _pvsOverride = default!;

    private ISawmill _sawmill = default!;

    /// <inheritdoc/>
    public override void Initialize()
    {
        base.Initialize();

        _sawmill = LogManager.GetSawmill("station");

        SubscribeLocalEvent<GameRunLevelChangedEvent>(OnRoundEnd);
        SubscribeLocalEvent<PostGameMapLoad>(OnPostGameMapLoad);
        SubscribeLocalEvent<StationDataComponent, ComponentStartup>(OnStationAdd);
        SubscribeLocalEvent<StationDataComponent, ComponentShutdown>(OnStationDeleted);
        SubscribeLocalEvent<StationMemberComponent, ComponentShutdown>(OnStationGridDeleted);
        SubscribeLocalEvent<StationMemberComponent, PostGridSplitEvent>(OnStationSplitEvent);

        SubscribeLocalEvent<StationGridAddedEvent>(OnStationGridAdded);
        SubscribeLocalEvent<StationGridRemovedEvent>(OnStationGridRemoved);

        _player.PlayerStatusChanged += OnPlayerStatusChanged;
    }

    private void OnStationSplitEvent(EntityUid uid, StationMemberComponent component, ref PostGridSplitEvent args)
    {
        AddGridToStation(component.Station, args.Grid); // Add the new grid as a member.
    }

    private void OnStationGridDeleted(EntityUid uid, StationMemberComponent component, ComponentShutdown args)
    {
        if (!TryComp<StationDataComponent>(component.Station, out var stationData))
            return;

        stationData.Grids.Remove(uid);
        Dirty(component.Station, stationData);
    }

    public override void Shutdown()
    {
        base.Shutdown();
        _player.PlayerStatusChanged -= OnPlayerStatusChanged;
    }

    private void OnPlayerStatusChanged(object? sender, SessionStatusEventArgs e)
    {
        if (e.NewStatus == SessionStatus.Connected)
        {
            RaiseNetworkEvent(new StationsUpdatedEvent(GetStationNames()), e.Session);
        }
    }

    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)
    {
        RaiseNetworkEvent(new StationsUpdatedEvent(GetStationNames()), Filter.Broadcast());

        var metaData = MetaData(uid);
        RaiseLocalEvent(new StationInitializedEvent(uid));
        _sawmill.Info($"Set up station {metaData.EntityName} ({uid}).");
        _pvsOverride.AddGlobalOverride(uid);
    }

    private void OnStationDeleted(EntityUid uid, StationDataComponent component, ComponentShutdown args)
    {
        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());
    }

    private void OnPostGameMapLoad(PostGameMapLoad ev)
    {
        var dict = new Dictionary<string, List<EntityUid>>();

        // Iterate over all BecomesStation
        foreach (var grid in ev.Grids)
        {
            // We still setup the grid
            if (TryComp<BecomesStationComponent>(grid, out var becomesStation))
                dict.GetOrNew(becomesStation.Id).Add(grid);
        }

        if (!dict.Any())
        {
            // Oh jeez, no stations got loaded.
            // We'll yell about it, but the thing this used to do with creating a dummy is kinda pointless now.
            _sawmill.Error($"There were no station grids for {ev.GameMap.ID}!");
        }

        foreach (var (id, gridIds) in dict)
        {
            StationConfig stationConfig;

            if (ev.GameMap.Stations.ContainsKey(id))
                stationConfig = ev.GameMap.Stations[id];
            else
            {
                _sawmill.Error($"The station {id} in map {ev.GameMap.ID} does not have an associated station config!");
                continue;
            }

            InitializeNewStation(stationConfig, gridIds, ev.StationName, ev.GameMap);
        }
    }

    private void OnRoundEnd(GameRunLevelChangedEvent eventArgs)
    {
        if (eventArgs.New != GameRunLevel.PreRoundLobby)
            return;

        var query = EntityQueryEnumerator<StationDataComponent>();
        while (query.MoveNext(out var station, out _))
        {
            QueueDel(station);
        }
    }

    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>
    /// Initializes a new station with the given information.
    /// </summary>
    /// <param name="stationConfig">The station configuration to use.</param>
    /// <param name="gridIds">All grids that should be added to the station.</param>
    /// <param name="name">Optional override for the station name.</param>
    /// <param name="gameMap">The game map that created this station, if any.</param>
    /// <remarks>This is for ease of use, manually spawning the entity works just fine.</remarks>
    /// <returns>The initialized station.</returns>
    public EntityUid InitializeNewStation(
        StationConfig stationConfig,
        IEnumerable<EntityUid>? gridIds,
        string? name = null,
        Content.Shared.Maps.GameMapPrototype? gameMap = null)
    {
        // Use overrides for setup.
        var station = EntityManager.SpawnEntity(stationConfig.StationPrototype, MapCoordinates.Nullspace, stationConfig.StationComponentOverrides);

        if (name is not null)
            RenameStation(station, name, false);

        DebugTools.Assert(HasComp<StationDataComponent>(station), "Stations should have StationData in their prototype.");

        var data = Comp<StationDataComponent>(station);
        data.JobWeights = gameMap?.JobWeights;
        Dirty(station, data);
        name ??= MetaData(station).EntityName;

        foreach (var grid in gridIds ?? Array.Empty<EntityUid>())
        {
            AddGridToStation(station, grid, null, data, name);
        }

        var ev = new StationPostInitEvent((station, data));
        RaiseLocalEvent(station, ref ev, true);

        return station;
    }

    /// <summary>
    /// Adds the given grid to a station.
    /// </summary>
    /// <param name="mapGrid">Grid to attach.</param>
    /// <param name="station">Station to attach the grid to.</param>
    /// <param name="gridComponent">Resolve pattern, grid component of mapGrid.</param>
    /// <param name="stationData">Resolve pattern, station data component of station.</param>
    /// <param name="name">The name to assign to the grid if any.</param>
    /// <exception cref="ArgumentException">Thrown when mapGrid or station are not a grid or station, respectively.</exception>
    public void AddGridToStation(EntityUid station, EntityUid mapGrid, MapGridComponent? gridComponent = null, StationDataComponent? stationData = null, string? name = null)
    {
        if (!Resolve(mapGrid, ref gridComponent))
            throw new ArgumentException("Tried to initialize a station on a non-grid entity!", nameof(mapGrid));
        if (!Resolve(station, ref stationData))
            throw new ArgumentException("Tried to use a non-station entity as a station!", nameof(station));

        if (!string.IsNullOrEmpty(name))
            Meta.SetEntityName(mapGrid, name);

        var stationMember = EnsureComp<StationMemberComponent>(mapGrid);
        stationMember.Station = station;
        stationData.Grids.Add(mapGrid);
        Dirty(station, stationData);
        Dirty(mapGrid, stationMember);

        RaiseLocalEvent(station, new StationGridAddedEvent(mapGrid, station, false), true);

        _sawmill.Info($"Adding grid {mapGrid} to station {Name(station)} ({station})");
    }

    /// <summary>
    /// Removes the given grid from a station.
    /// </summary>
    /// <param name="station">Station to remove the grid from.</param>
    /// <param name="mapGrid">Grid to remove</param>
    /// <param name="gridComponent">Resolve pattern, grid component of mapGrid.</param>
    /// <param name="stationData">Resolve pattern, station data component of station.</param>
    /// <exception cref="ArgumentException">Thrown when mapGrid or station are not a grid or station, respectively.</exception>
    public void RemoveGridFromStation(EntityUid station, EntityUid mapGrid, MapGridComponent? gridComponent = null, StationDataComponent? stationData = null)
    {
        if (!Resolve(mapGrid, ref gridComponent))
            throw new ArgumentException("Tried to initialize a station on a non-grid entity!", nameof(mapGrid));
        if (!Resolve(station, ref stationData))
            throw new ArgumentException("Tried to use a non-station entity as a station!", nameof(station));

        RemComp<StationMemberComponent>(mapGrid);
        stationData.Grids.Remove(mapGrid);
        Dirty(station, stationData);

        RaiseLocalEvent(station, new StationGridRemovedEvent(mapGrid, station), true);
        _sawmill.Info($"Removing grid {mapGrid} from station {Name(station)} ({station})");
    }

    /// <summary>
    /// Renames the given station.
    /// </summary>
    /// <param name="station">Station to rename.</param>
    /// <param name="name">The new name to apply.</param>
    /// <param name="loud">Whether or not to announce the rename.</param>
    /// <param name="stationData">Resolve pattern, station data component of station.</param>
    /// <param name="metaData">Resolve pattern, metadata component of station.</param>
    /// <exception cref="ArgumentException">Thrown when the given station is not a station.</exception>
    public void RenameStation(EntityUid station, string name, bool loud = true, StationDataComponent? stationData = null, MetaDataComponent? metaData = null)
    {
        if (!Resolve(station, ref stationData, ref metaData))
            throw new ArgumentException("Tried to use a non-station entity as a station!", nameof(station));

        var oldName = metaData.EntityName;
        Meta.SetEntityName(station, name, metaData);

        if (loud)
        {
            _chatSystem.DispatchStationAnnouncement(station, $"The station {oldName} has been renamed to {name}.");
        }

        RaiseLocalEvent(station, new StationRenamedEvent(oldName, name), true);
    }

    /// <summary>
    /// Deletes the given station.
    /// </summary>
    /// <param name="station">Station to delete.</param>
    /// <param name="stationData">Resolve pattern, station data component of station.</param>
    /// <exception cref="ArgumentException">Thrown when the given station is not a station.</exception>
    public void DeleteStation(EntityUid station, StationDataComponent? stationData = null)
    {
        if (!Resolve(station, ref stationData))
            throw new ArgumentException("Tried to use a non-station entity as a station!", nameof(station));

        QueueDel(station);
    }

    /// <inheritdoc/>
    public override bool TryFindRandomTileOnStation(Entity<StationDataComponent> station,
        out Vector2i tile,
        [NotNullWhen(true)] out Entity<MapGridComponent>? targetGrid,
        out EntityCoordinates targetCoords,
        int numAttempts = 10)
    {
        tile = default;
        targetCoords = EntityCoordinates.Invalid;
        targetGrid = null;

        // Weight grid choice by tilecount
        var totalTiles = 0;
        var grids = new List<(Entity<MapGridComponent> Entity, int Count, List<TileRef> Tiles)>();
        foreach (var possibleTarget in station.Comp.Grids)
        {
            if (!GridQuery.TryComp(possibleTarget, out var comp))
                continue;

            // Get the tile count for the given grid.
            var tileCount = Map.GetFilledTileCount((possibleTarget, comp));

            // Just to be sure, no empty elements.
            if (tileCount > 0)
            {
                grids.Add(((possibleTarget, comp), tileCount, new()));
                totalTiles += tileCount;
            }
        }

        if (grids.Count == 0)
            return false;

        for (var i = 0; i < numAttempts; i++)
        {
            // Find random tile within list.
            var nextTileIndex = Random.Next(totalTiles);
            TileRef? randomTileRef = null;
            var startIndex = 0;
            for (int j = 0; j < grids.Count; j++)
            {
                var grid = grids[j];
                // If the index is in this particular grid, find it and remove the tile to prevent selecting it twice.
                if (nextTileIndex >= startIndex + grid.Count)
                {
                    startIndex += grid.Count;
                    continue;
                }

                targetGrid = grid.Entity;

                // Empty list: hasn't been queried yet - get our tiles.
                if (grid.Tiles.Count <= 0)
                {
                    grid.Tiles = Map.GetAllTiles(targetGrid.Value, targetGrid.Value.Comp).ToList();

                    // Actual list count doesn't match expected count (a bug - return failure).
                    Debug.Assert(grid.Tiles.Count == grid.Count);
                    if (grid.Tiles.Count != grid.Count)
                        return false;
                }

                var ourTileIndex = nextTileIndex - startIndex;
                randomTileRef = grid.Tiles[ourTileIndex];
                grid.Tiles.RemoveSwap(ourTileIndex);
                grid.Count--;
                totalTiles--;

                // Empty list, remove element
                if (grid.Tiles.Count <= 0)
                    grids.RemoveSwap(j);

                break;
            }

            // Out of valid tiles, return early.
            if (randomTileRef is not { } tileRef)
                return false;

            // Could not find a grid...
            if (targetGrid == null)
                continue;

            // Invalid tile, try again.
            if (_atmos.IsTileSpace(targetGrid.Value.Owner, Transform(targetGrid.Value).MapUid, tileRef.GridIndices)
                || _atmos.IsTileAirBlockedCached(targetGrid.Value.Owner, tileRef.GridIndices))
            {
                continue;
            }

            targetCoords = Map.GridTileToLocal(targetGrid.Value, targetGrid.Value, tileRef.GridIndices);
            tile = tileRef.GridIndices;
            return true;
        }

        return false;
    }
}

/// <summary>
/// Broadcast event fired when a station is first set up.
/// This is the ideal point to add components to it.
/// </summary>
[PublicAPI]
public sealed class StationInitializedEvent : EntityEventArgs
{
    /// <summary>
    /// Station this event is for.
    /// </summary>
    public EntityUid Station;

    public StationInitializedEvent(EntityUid station)
    {
        Station = station;
    }
}

/// <summary>
/// Directed event fired on a station when a grid becomes a member of the station.
/// </summary>
[PublicAPI]
public sealed class StationGridAddedEvent : EntityEventArgs
{
    /// <summary>
    /// ID of the grid added to the station.
    /// </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, EntityUid station, bool isSetup)
    {
        GridId = gridId;
        Station = station;
        IsSetup = isSetup;
    }
}

/// <summary>
/// Directed event fired on a station when a grid is no longer a member of the station.
/// </summary>
[PublicAPI]
public sealed class StationGridRemovedEvent : EntityEventArgs
{
    /// <summary>
    /// ID of the grid removed from the station.
    /// </summary>
    public 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;
    }
}

/// <summary>
/// Directed event fired on a station when it is renamed.
/// </summary>
[PublicAPI]
public sealed class StationRenamedEvent : EntityEventArgs
{
    /// <summary>
    /// Prior name of the station.
    /// </summary>
    public string OldName;

    /// <summary>
    /// New name of the station.
    /// </summary>
    public string NewName;

    public StationRenamedEvent(string oldName, string newName)
    {
        OldName = oldName;
        NewName = newName;
    }
}