tbd-station-14/Content.Shared/DeviceLinking/SharedDeviceLinkSystem.cs

using Content.Shared.Administration.Logs;
using Content.Shared.Database;
using Content.Shared.DeviceLinking.Events;
using Content.Shared.Popups;
using Robust.Shared.Prototypes;
using Robust.Shared.Utility;

namespace Content.Shared.DeviceLinking;

public abstract class SharedDeviceLinkSystem : EntitySystem
{
    [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
    [Dependency] private readonly SharedPopupSystem _popupSystem = default!;
    [Dependency] private readonly ISharedAdminLogManager _adminLogger = default!;
    private ISawmill _sawmill = default!;

    public const string InvokedPort = "link_port";

    /// <inheritdoc/>
    public override void Initialize()
    {
        SubscribeLocalEvent<DeviceLinkSourceComponent, ComponentInit>(OnInit);
        SubscribeLocalEvent<DeviceLinkSourceComponent, ComponentStartup>(OnSourceStartup);
        SubscribeLocalEvent<DeviceLinkSinkComponent, ComponentStartup>(OnSinkStartup);
        SubscribeLocalEvent<DeviceLinkSourceComponent, ComponentRemove>(OnSourceRemoved);
        SubscribeLocalEvent<DeviceLinkSinkComponent, ComponentRemove>(OnSinkRemoved);
        _sawmill = Logger.GetSawmill("devicelink");
    }

    #region Link Validation

    private void OnInit(EntityUid uid, DeviceLinkSourceComponent component, ComponentInit args)
    {
        // Populate the output dictionary.
        foreach (var (sinkUid, links) in component.LinkedPorts)
        {
            foreach (var link in links)
            {
                component.Outputs.GetOrNew(link.source).Add(sinkUid);
            }
        }
    }

    /// <summary>
    /// Removes invalid links where the saved sink doesn't exist/have a sink component for example
    /// </summary>
    private void OnSourceStartup(EntityUid sourceUid, DeviceLinkSourceComponent sourceComponent, ComponentStartup args)
    {
        List<EntityUid> invalidSinks = new();
        foreach (var sinkUid  in sourceComponent.LinkedPorts.Keys)
        {
            if (!TryComp<DeviceLinkSinkComponent>(sinkUid, out var sinkComponent))
            {
                invalidSinks.Add(sinkUid);
                foreach (var savedSinks in sourceComponent.Outputs.Values)
                {
                    savedSinks.Remove(sinkUid);
                }

                continue;
            }

            sinkComponent.LinkedSources.Add(sourceUid);
        }

        foreach (var invalidSink in invalidSinks)
        {
            sourceComponent.LinkedPorts.Remove(invalidSink);
        }
    }

    /// <summary>
    /// Same with <see cref="OnSourceStartup"/> but also checks that the saved ports are present on the sink
    /// </summary>
    private void OnSinkStartup(EntityUid sinkUid, DeviceLinkSinkComponent sinkComponent, ComponentStartup args)
    {
        List<EntityUid> invalidSources = new();
        foreach (var sourceUid in sinkComponent.LinkedSources)
        {
            if (!TryComp<DeviceLinkSourceComponent>(sourceUid, out var sourceComponent))
            {
                invalidSources.Add(sourceUid);
                continue;
            }

            if (!sourceComponent.LinkedPorts.TryGetValue(sinkUid, out var linkedPorts))
            {
                foreach (var savedSinks in sourceComponent.Outputs.Values)
                {
                    savedSinks.Remove(sinkUid);
                }
                continue;
            }

            if (sinkComponent.Ports == null)
                continue;

            List<(string, string)> invalidLinks = new();
            foreach (var link in linkedPorts)
            {
                if (!sinkComponent.Ports.Contains(link.sink))
                    invalidLinks.Add(link);
            }

            foreach (var invalidLink in invalidLinks)
            {
                linkedPorts.Remove(invalidLink);
                sourceComponent.Outputs.GetValueOrDefault(invalidLink.Item1)?.Remove(sinkUid);
            }
        }

        foreach (var invalidSource in invalidSources)
        {
            sinkComponent.LinkedSources.Remove(invalidSource);
        }
    }
    #endregion

    /// <summary>
    /// Ensures that its links get deleted when a source gets removed
    /// </summary>
    private void OnSourceRemoved(EntityUid uid, DeviceLinkSourceComponent component, ComponentRemove args)
    {
        var query = GetEntityQuery<DeviceLinkSinkComponent>();
        foreach (var sinkUid in component.LinkedPorts.Keys)
        {
            if (query.TryGetComponent(sinkUid, out var sink))
                RemoveSinkFromSourceInternal(uid, sinkUid, component, sink);
        }
    }

    /// <summary>
    /// Ensures that its links get deleted when a sink gets removed
    /// </summary>
    private void OnSinkRemoved(EntityUid sinkUid, DeviceLinkSinkComponent sinkComponent, ComponentRemove args)
    {
        var query = GetEntityQuery<DeviceLinkSourceComponent>();
        foreach (var linkedSource in sinkComponent.LinkedSources)
        {
            if (query.TryGetComponent(sinkUid, out var source))
                RemoveSinkFromSourceInternal(linkedSource, sinkUid, source, sinkComponent);
        }
    }

    #region Ports
    /// <summary>
    /// Convenience function to add several ports to an entity
    /// </summary>
    public void EnsureSourcePorts(EntityUid uid, params string[] ports)
    {
        var comp = EnsureComp<DeviceLinkSourceComponent>(uid);
        comp.Ports ??= new HashSet<string>();

        foreach (var port in ports)
        {
            DebugTools.Assert(_prototypeManager.HasIndex<SourcePortPrototype>(port));
            comp.Ports?.Add(port);
        }
    }

    /// <summary>
    /// Convenience function to add several ports to an entity.
    /// </summary>
    public void EnsureSinkPorts(EntityUid uid, params string[] ports)
    {
        var comp = EnsureComp<DeviceLinkSinkComponent>(uid);
        comp.Ports ??= new HashSet<string>();

        foreach (var port in ports)
        {
            DebugTools.Assert(_prototypeManager.HasIndex<SinkPortPrototype>(port));
            comp.Ports?.Add(port);
        }
    }

    /// <summary>
    /// Retrieves the available ports from a source
    /// </summary>
    /// <returns>A list of source port prototypes</returns>
    public List<SourcePortPrototype> GetSourcePorts(EntityUid sourceUid, DeviceLinkSourceComponent? sourceComponent = null)
    {
        if (!Resolve(sourceUid, ref sourceComponent) || sourceComponent.Ports == null)
            return new List<SourcePortPrototype>();

        var sourcePorts = new List<SourcePortPrototype>();
        foreach (var port in sourceComponent.Ports)
        {
            sourcePorts.Add(_prototypeManager.Index<SourcePortPrototype>(port));
        }

        return sourcePorts;
    }

    /// <summary>
    /// Retrieves the available ports from a sink
    /// </summary>
    /// <returns>A list of sink port prototypes</returns>
    public List<SinkPortPrototype> GetSinkPorts(EntityUid sinkUid, DeviceLinkSinkComponent? sinkComponent = null)
    {
        if (!Resolve(sinkUid, ref sinkComponent) || sinkComponent.Ports == null)
            return new List<SinkPortPrototype>();

        var sinkPorts = new List<SinkPortPrototype>();
        foreach (var port in sinkComponent.Ports)
        {
            sinkPorts.Add(_prototypeManager.Index<SinkPortPrototype>(port));
        }

        return sinkPorts;
    }

    /// <summary>
    /// Convenience function to retrieve the name of a port prototype
    /// </summary>
    public string PortName<TPort>(string port) where TPort : DevicePortPrototype, IPrototype
    {
        if (!_prototypeManager.TryIndex<TPort>(port, out var proto))
            return port;

        return Loc.GetString(proto.Name);
    }
    #endregion

    #region Links
    /// <summary>
    /// Returns the links of a source
    /// </summary>
    /// <returns>A list of sink and source port ids that are linked together</returns>
    public HashSet<(string source, string sink)> GetLinks(EntityUid sourceUid, EntityUid sinkUid, DeviceLinkSourceComponent? sourceComponent = null)
    {
        if (!Resolve(sourceUid, ref sourceComponent) || !sourceComponent.LinkedPorts.TryGetValue(sinkUid, out var links))
            return new HashSet<(string source, string sink)>();

        return links;
    }

    /// <summary>
    /// Returns the default links for the given list of source port prototypes
    /// </summary>
    /// <param name="sources">The list of source port prototypes to get the default links for</param>
    /// <returns>A list of sink and source port ids</returns>
    public List<(string source, string sink)> GetDefaults(List<SourcePortPrototype> sources)
    {
        var defaults = new List<(string, string)>();
        foreach (var source in sources)
        {
            if (source.DefaultLinks == null)
                return new List<(string, string)>();

            foreach (var defaultLink in source.DefaultLinks)
            {
                defaults.Add((source.ID, defaultLink));
            }
        }

        return defaults;
    }

    /// <summary>
    /// Links the given source and sink by their default links
    /// </summary>
    /// <param name="userId">Optinal user uid for displaying popups</param>
    /// <param name="sourceUid">The source uid</param>
    /// <param name="sinkUid">The sink uid</param>
    /// <param name="sourceComponent"></param>
    /// <param name="sinkComponent"></param>
    public void LinkDefaults(
        EntityUid? userId,
        EntityUid sourceUid,
        EntityUid sinkUid,
        DeviceLinkSourceComponent? sourceComponent = null,
        DeviceLinkSinkComponent? sinkComponent = null)
    {
        if (!Resolve(sourceUid, ref sourceComponent) || !Resolve(sinkUid, ref sinkComponent))
            return;

        if (userId != null)
            _adminLogger.Add(LogType.DeviceLinking, LogImpact.Low, $"{ToPrettyString(userId.Value):actor} is linking defaults between {ToPrettyString(sourceUid):source} and {ToPrettyString(sinkUid):sink}");
        else
            _adminLogger.Add(LogType.DeviceLinking, LogImpact.Low, $"linking defaults between {ToPrettyString(sourceUid):source} and {ToPrettyString(sinkUid):sink}");

        var sourcePorts = GetSourcePorts(sourceUid, sourceComponent);
        var defaults = GetDefaults(sourcePorts);
        SaveLinks(userId, sourceUid, sinkUid, defaults, sourceComponent, sinkComponent);

        if (userId != null)
            _popupSystem.PopupCursor(Loc.GetString("signal-linking-verb-success", ("machine", sourceUid)), userId.Value);
    }


    /// <summary>
    /// Saves multiple links between a source and a sink device.
    /// Ignores links where either the source or sink port aren't present
    /// </summary>
    /// <param name="userId">Optinal user uid for displaying popups</param>
    /// <param name="sourceUid">The source uid</param>
    /// <param name="sinkUid">The sink uid</param>
    /// <param name="links">List of source and sink ids to link</param>
    /// <param name="sourceComponent"></param>
    /// <param name="sinkComponent"></param>
    public void SaveLinks(
        EntityUid? userId,
        EntityUid sourceUid,
        EntityUid sinkUid,
        List<(string source, string sink)> links,
        DeviceLinkSourceComponent? sourceComponent = null,
        DeviceLinkSinkComponent? sinkComponent = null)
    {
        if (!Resolve(sourceUid, ref sourceComponent) || !Resolve(sinkUid, ref sinkComponent))
            return;

        if (sourceComponent.Ports == null || sinkComponent.Ports == null)
            return;

        if (!InRange(sourceUid, sinkUid, sourceComponent.Range))
        {
            if (userId != null)
                _popupSystem.PopupCursor(Loc.GetString("signal-linker-component-out-of-range"), userId.Value);

            return;
        }

        RemoveSinkFromSource(sourceUid, sinkUid, sourceComponent);
        foreach (var (source, sink) in links)
        {
            DebugTools.Assert(_prototypeManager.HasIndex<SourcePortPrototype>(source));
            DebugTools.Assert(_prototypeManager.HasIndex<SinkPortPrototype>(sink));

            if (!sourceComponent.Ports.Contains(source) || !sinkComponent.Ports.Contains(sink))
                continue;

            if (!CanLink(userId, sourceUid, sinkUid, source, sink, false, sourceComponent))
                continue;

            sourceComponent.Outputs.GetOrNew(source).Add(sinkUid);
            sourceComponent.LinkedPorts.GetOrNew(sinkUid).Add((source, sink));

            SendNewLinkEvent(userId, sourceUid, source, sinkUid, sink);
        }