CatStation/Content.Shared/Interaction/SmartEquipSystem.cs

using Content.Shared.ActionBlocker;
using Content.Shared.Containers.ItemSlots;
using Content.Shared.Hands.Components;
using Content.Shared.Hands.EntitySystems;
using Content.Shared.Input;
using Content.Shared.Inventory;
using Content.Shared.Popups;
using Content.Shared.Stacks;
using Content.Shared.Storage;
using Content.Shared.Storage.EntitySystems;
using Content.Shared.Whitelist;
using Robust.Shared.Containers;
using Robust.Shared.Input.Binding;
using Robust.Shared.Player;

namespace Content.Shared.Interaction;

/// <summary>
/// This handles smart equipping or inserting/ejecting from slots through keybinds--generally shift+E and shift+B
/// </summary>
public sealed partial class SmartEquipSystem : EntitySystem
{
    [Dependency] private SharedHandsSystem _hands = default!;
    [Dependency] private SharedStorageSystem _storage = default!;
    [Dependency] private InventorySystem _inventory = default!;
    [Dependency] private ItemSlotsSystem _slots = default!;
    [Dependency] private SharedContainerSystem _container = default!;
    [Dependency] private SharedPopupSystem _popup = default!;
    [Dependency] private ActionBlockerSystem _actionBlocker = default!;
    [Dependency] private EntityWhitelistSystem _whitelistSystem = default!;

    /// <inheritdoc/>
    public override void Initialize()
    {
        CommandBinds.Builder
            .Bind(ContentKeyFunctions.SmartEquipBackpack, InputCmdHandler.FromDelegate(HandleSmartEquipBackpack, handle: false, outsidePrediction: false))
            .Bind(ContentKeyFunctions.SmartEquipBelt, InputCmdHandler.FromDelegate(HandleSmartEquipBelt, handle: false, outsidePrediction: false))
            .Bind(ContentKeyFunctions.SmartEquipPocket1, InputCmdHandler.FromDelegate(HandleSmartEquipPocket1, handle: false, outsidePrediction: false))
            .Bind(ContentKeyFunctions.SmartEquipPocket2, InputCmdHandler.FromDelegate(HandleSmartEquipPocket2, handle: false, outsidePrediction: false))
            .Bind(ContentKeyFunctions.SmartEquipSuitStorage, InputCmdHandler.FromDelegate(HandleSmartEquipSuitStorage, handle: false, outsidePrediction: false))
            .Register<SmartEquipSystem>();
    }

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

        CommandBinds.Unregister<SmartEquipSystem>();
    }

    private void HandleSmartEquipBackpack(ICommonSession? session)
    {
        HandleSmartEquip(session, "back");
    }

    private void HandleSmartEquipBelt(ICommonSession? session)
    {
        HandleSmartEquip(session, "belt");
    }

    private void HandleSmartEquipPocket1(ICommonSession? session)
    {
        HandleSmartEquip(session, "pocket1");
    }

    private void HandleSmartEquipPocket2(ICommonSession? session)
    {
        HandleSmartEquip(session, "pocket2");
    }

    private void HandleSmartEquipSuitStorage(ICommonSession? session)
    {
        HandleSmartEquip(session, "suitstorage");
    }

    private void HandleSmartEquip(ICommonSession? session, string equipmentSlot)
    {
        if (session is not { } playerSession)
            return;

        if (playerSession.AttachedEntity is not { Valid: true } uid || !Exists(uid))
            return;

        // early out if we don't have any hands or a valid inventory slot
        if (!TryComp<HandsComponent>(uid, out var hands) || hands.ActiveHandId == null)
            return;

        var handItem = _hands.GetActiveItem((uid, hands));

        // can the user interact, and is the item interactable? e.g. virtual items
        if (!_actionBlocker.CanInteract(uid, handItem))
            return;

        if (!TryComp<InventoryComponent>(uid, out var inventory) || !_inventory.HasSlot(uid, equipmentSlot, inventory))
        {
            _popup.PopupEntity(Loc.GetString("smart-equip-missing-equipment-slot", ("slotName", equipmentSlot)), uid, uid);
            return;
        }

        // early out if we have an item and cant drop it at all
        if (handItem != null && !_hands.CanDropHeld(uid, hands.ActiveHandId))
        {
            _popup.PopupEntity(Loc.GetString("smart-equip-cant-drop"), uid, uid);
            return;
        }

        // There are eight main cases we want to handle here,
        // so let's write them out

        // if the slot we're trying to smart equip from:
        // 1) doesn't have an item
        //    - with hand item: try to put it in the slot
        //    - without hand item: fail
        // 2) has an item, and that item is a storage item
        //    - with hand item: try to put it in storage
        //    - without hand item: try to take the last stored item and put it in our hands
        // 3) has an item, and that item is an item slots holder
        //    - with hand item: get the highest priority item slot with a valid whitelist and try to insert it
        //    - without hand item: get the highest priority item slot with an item and try to eject it
        // 4) has an item, with no special storage components
        //    - with hand item: fail
        //    - without hand item: try to put the item into your hand

        _inventory.TryGetSlotEntity(uid, equipmentSlot, out var slotEntity);
        var emptyEquipmentSlotString = Loc.GetString("smart-equip-empty-equipment-slot", ("slotName", equipmentSlot));

        // case 1 (no slot item):
        if (slotEntity is not { } slotItem)
        {
            if (handItem == null)
            {
                _popup.PopupEntity(emptyEquipmentSlotString, uid, uid);
                return;
            }

            if (!_inventory.CanEquip(uid, handItem.Value, equipmentSlot, out var reason))
            {
                _popup.PopupEntity(Loc.GetString(reason), uid, uid);
                return;
            }

            _hands.TryDrop((uid, hands), hands.ActiveHandId!);
            _inventory.TryEquip(uid, handItem.Value, equipmentSlot, predicted: true, checkDoafter:true);
            return;
        }

        // case 2 (storage item):
        if (TryComp<StorageComponent>(slotItem, out var storage))
        {
            switch (handItem)
            {
                case null when storage.Container.ContainedEntities.Count == 0:
                    _popup.PopupEntity(emptyEquipmentSlotString, uid, uid);
                    return;
                case null:
                    var removing = storage.Container.ContainedEntities[^1];
                    _container.RemoveEntity(slotItem, removing);
                    _hands.TryPickup(uid, removing, handsComp: hands);
                    return;
            }

            if (!_storage.CanInsert(slotItem, handItem.Value, out var reason))
            {
                if (reason != null)
                    _popup.PopupEntity(Loc.GetString(reason), uid, uid);

                return;
            }

            _hands.TryDrop((uid, hands), hands.ActiveHandId!);
            _storage.Insert(slotItem, handItem.Value, out var stacked, out _, user: uid);

            // if the hand item stacked with the things in inventory, but there's no more space left for the rest
            // of the stack, place the stack back in hand rather than dropping it on the floor
            if (stacked != null && !_storage.CanInsert(slotItem, handItem.Value, out _))
            {
                if (TryComp<StackComponent>(handItem.Value, out var handStack) && handStack.Count > 0)
                    _hands.TryPickup(uid, handItem.Value, handsComp: hands);
            }

            return;
        }

        // case 3 (itemslot item):
        if (TryComp<ItemSlotsComponent>(slotItem, out var slots))
        {
            if (handItem == null)
            {
                ItemSlot? toEjectFrom = null;

                foreach (var slot in slots.Slots.Values)
                {
                    if (slot.HasItem && slot.Priority > (toEjectFrom?.Priority ?? int.MinValue))
                        toEjectFrom = slot;
                }

                if (toEjectFrom == null)
                {
                    _popup.PopupEntity(emptyEquipmentSlotString, uid, uid);
                    return;
                }

                _slots.TryEjectToHands(slotItem, toEjectFrom, uid, excludeUserAudio: true);
                return;
            }

            ItemSlot? toInsertTo = null;

            foreach (var slot in slots.Slots.Values)
            {
                if (!slot.HasItem
                    && _whitelistSystem.IsWhitelistPassOrNull(slot.Whitelist, handItem.Value)
                    && slot.Priority > (toInsertTo?.Priority ?? int.MinValue))
                {
                    toInsertTo = slot;
                }
            }

            if (toInsertTo == null)
            {
                _popup.PopupEntity(Loc.GetString("smart-equip-no-valid-item-slot-insert", ("item", handItem.Value)), uid, uid);
                return;
            }

            _slots.TryInsertFromHand(slotItem, toInsertTo, uid, hands, excludeUserAudio: true);
            return;
        }

        // case 4 (just an item):
        if (handItem != null)
            return;

        if (!_inventory.CanUnequip(uid, equipmentSlot, out var inventoryReason))
        {
            _popup.PopupEntity(Loc.GetString(inventoryReason), uid, uid);
            return;
        }

        _inventory.TryUnequip(uid, equipmentSlot, inventory: inventory, predicted: true, checkDoafter: true);
        _hands.TryPickup(uid, slotItem, handsComp: hands);
    }
}