CatStation/Content.Shared/Body/BodySystem.cs

using Content.Shared.DragDrop;
using JetBrains.Annotations;
using Robust.Shared.Containers;

namespace Content.Shared.Body;

/// <summary>
/// System responsible for coordinating entities with <see cref="BodyComponent" /> and their entities with <see cref="OrganComponent" />.
/// This system is primarily responsible for event relaying and the relationships between a body and its organs.
/// It is not responsible for player-facing body features, e.g. "blood" or "breathing."
/// Such features should be implemented in systems relying on the various events raised by this class.
/// </summary>
/// <seealso cref="OrganGotInsertedEvent" />
/// <seealso cref="OrganGotRemovedEvent" />
/// <seealso cref="OrganInsertedIntoEvent" />
/// <seealso cref="OrganRemovedFromEvent" />
/// <seealso cref="BodyRelayedEvent{TEvent}" />
public sealed partial class BodySystem : EntitySystem
{
    [Dependency] private SharedContainerSystem _container = default!;

    [Dependency] private EntityQuery<BodyComponent> _bodyQuery = default!;
    [Dependency] private EntityQuery<OrganComponent> _organQuery = default!;

    public override void Initialize()
    {
        base.Initialize();

        SubscribeLocalEvent<BodyComponent, ComponentInit>(OnBodyInit);
        SubscribeLocalEvent<BodyComponent, ComponentShutdown>(OnBodyShutdown);

        SubscribeLocalEvent<BodyComponent, CanDragEvent>(OnCanDrag);

        SubscribeLocalEvent<BodyComponent, EntInsertedIntoContainerMessage>(OnBodyEntInserted);
        SubscribeLocalEvent<BodyComponent, EntRemovedFromContainerMessage>(OnBodyEntRemoved);

        InitializeRelay();
    }

    private void OnBodyInit(Entity<BodyComponent> ent, ref ComponentInit args)
    {
        ent.Comp.Organs =
            _container.EnsureContainer<Container>(ent, BodyComponent.ContainerID);
    }

    private void OnBodyShutdown(Entity<BodyComponent> ent, ref ComponentShutdown args)
    {
        if (ent.Comp.Organs is { } organs)
            _container.ShutdownContainer(organs);
    }

    private void OnBodyEntInserted(Entity<BodyComponent> ent, ref EntInsertedIntoContainerMessage args)
    {
        if (args.Container.ID != BodyComponent.ContainerID)
            return;

        if (!_organQuery.TryComp(args.Entity, out var organ))
            return;

        var body = new OrganInsertedIntoEvent(args.Entity);
        RaiseLocalEvent(ent, ref body);

        var ev = new OrganGotInsertedEvent(ent);
        RaiseLocalEvent(args.Entity, ref ev);

        if (organ.Body != ent)
        {
            organ.Body = ent;
            Dirty(args.Entity, organ);
        }
    }

    private void OnBodyEntRemoved(Entity<BodyComponent> ent, ref EntRemovedFromContainerMessage args)
    {
        if (args.Container.ID != BodyComponent.ContainerID)
            return;

        if (!_organQuery.TryComp(args.Entity, out var organ))
            return;

        var body = new OrganRemovedFromEvent(args.Entity);
        RaiseLocalEvent(ent, ref body);

        var ev = new OrganGotRemovedEvent(ent);
        RaiseLocalEvent(args.Entity, ref ev);

        if (organ.Body == null)
            return;

        organ.Body = null;
        Dirty(args.Entity, organ);
    }

    private void OnCanDrag(Entity<BodyComponent> ent, ref CanDragEvent args)
    {
        args.Handled = true;
    }

    /// <summary>
    /// Gets an enumerator of organs with a specific component.
    /// </summary>
    /// <param name="body">The entity to enumerate the organs of.</param>
    /// <typeparam name="T">The component.</typeparam>
    /// <returns>The enumerator of the entity's organs with the specified component.</returns>
    [PublicAPI]
    public IEnumerable<Entity<OrganComponent,T>> EnumerateOrgans<T>(Entity<BodyComponent?> body) where T : IComponent
    {
        if (!Resolve(body, ref body.Comp, false))
            yield break;

        foreach (var organ in body.Comp.Organs?.ContainedEntities ?? [])
        {
            if (TryComp<T>(organ, out var comp))
                yield return (organ, _organQuery.Comp(organ), comp);
        }
    }

    /// <summary>
    /// Gets an enumerator of organs with a specific component using a <see cref="EntityQuery{TComp1}"/>.
    /// </summary>
    /// <param name="body">The entity to enumerate the organs of.</param>
    /// <param name="query">The EntityQuery to use.</param>
    /// <typeparam name="T">The component.</typeparam>
    /// <returns>The enumerator of the entity's organs with the specified component.</returns>
    [PublicAPI]
    public IEnumerable<Entity<OrganComponent, T>> EnumerateOrgans<T>(Entity<BodyComponent?> body, EntityQuery<T> query) where T : IComponent
    {
        if (!Resolve(body, ref body.Comp, false))
            yield break;

        foreach (var organ in body.Comp.Organs?.ContainedEntities ?? [])
        {
            if (query.TryComp(organ, out var comp))
                yield return (organ, _organQuery.Comp(organ), comp);
        }
    }
}