Project-M/Assets/_Project/Scripts/Server/Economy/EquipSystem.cs

using ProjectM.Simulation;
using Unity.Burst;
using Unity.Collections;
using Unity.Entities;
using Unity.NetCode;

namespace ProjectM.Server
{
    /// <summary>
    /// Server-authoritative equipment handler (<see cref="EquipRequest"/> / <see cref="UnequipRequest"/> RPCs).
    /// Resolves the sender's player (SourceConnection -&gt; NetworkId -&gt; GhostOwner, the AbilityUpgradeSystem /
    /// InventoryDepositSystem owner-map idiom) and applies the change IN-PLACE: moves the item between the
    /// personal <see cref="InventorySlot"/> bag and the <see cref="EquipmentSlot"/> loadout (buffer index = slot),
    /// sets <see cref="AbilityRef"/>.Id from the Weapon slot (restoring <see cref="DefaultAbility"/> on
    /// weapon-unequip), and adds/strips the item's inline stat mods as <see cref="StatModifier"/>s tagged by a
    /// per-slot SourceId (<c>Tuning.EquipSourceIdBase + slot</c>), stripped TARGET-AGNOSTICALLY via
    /// <see cref="TimedModifierUtil.RemoveBySourceId"/>.
    ///
    /// Effects are EVENT-DRIVEN (applied once here): AbilityRef + StatModifier are [GhostField]s re-folded by the
    /// predicted StatRecomputeSystem every tick and replicated to the owner, so the swap is prediction-correct
    /// (DebugModifierInjectionSystem.CycleAbility is the precedent) and survives respawn (the entity persists).
    /// Atomicity: an equip into an occupied slot verifies the bag can hold the swapped-out item BEFORE any
    /// withdrawal and rejects otherwise — no item loss (the co-op-placement commit-in-place rule). Plain server
    /// SimulationSystemGroup (NOT predicted -&gt; applied once, no rollback double-apply); only the request entity
    /// destroy is deferred to the ECB.
    /// </summary>
    [BurstCompile]
    [WorldSystemFilter(WorldSystemFilterFlags.ServerSimulation)]
    public partial struct EquipSystem : ISystem
    {
        [BurstCompile]
        public void OnCreate(ref SystemState state)
        {
            state.RequireForUpdate<ItemDatabase>();
            var builder = new EntityQueryBuilder(Allocator.Temp)
                .WithAny<EquipRequest, UnequipRequest>().WithAll<ReceiveRpcCommandRequest>();
            state.RequireForUpdate(state.GetEntityQuery(builder));
        }

        [BurstCompile]
        public void OnUpdate(ref SystemState state)
        {
            var itemDb = SystemAPI.GetSingleton<ItemDatabase>();

            var playerByConn = new NativeHashMap<int, Entity>(8, Allocator.Temp);
            foreach (var (owner, entity) in
                     SystemAPI.Query<RefRO<GhostOwner>>()
                         .WithAll<PlayerTag, InventorySlot, EquipmentSlot>().WithEntityAccess())
                playerByConn[owner.ValueRO.NetworkId] = entity;

            var ecb = new EntityCommandBuffer(Allocator.Temp);

            foreach (var (request, receive, requestEntity) in
                     SystemAPI.Query<RefRO<EquipRequest>, RefRO<ReceiveRpcCommandRequest>>().WithEntityAccess())
            {
                if (TryResolvePlayer(ref state, playerByConn, receive.ValueRO.SourceConnection, out var player))
                    HandleEquip(ref state, itemDb, player, request.ValueRO.ItemId);
                ecb.DestroyEntity(requestEntity);
            }

            foreach (var (request, receive, requestEntity) in
                     SystemAPI.Query<RefRO<UnequipRequest>, RefRO<ReceiveRpcCommandRequest>>().WithEntityAccess())
            {
                if (TryResolvePlayer(ref state, playerByConn, receive.ValueRO.SourceConnection, out var player))
                    HandleUnequip(ref state, itemDb, player, request.ValueRO.Slot);
                ecb.DestroyEntity(requestEntity);
            }

            ecb.Playback(state.EntityManager);
            playerByConn.Dispose();
        }

        static bool TryResolvePlayer(ref SystemState state, NativeHashMap<int, Entity> map, Entity conn, out Entity player)
        {
            player = Entity.Null;
            return state.EntityManager.HasComponent<NetworkId>(conn)
                && map.TryGetValue(state.EntityManager.GetComponentData<NetworkId>(conn).Value, out player);
        }

        static void HandleEquip(ref SystemState state, ItemDatabase itemDb, Entity player, ushort itemId)
        {
            ref var db = ref itemDb.Value.Value;
            if (!db.TryGetItem(itemId, out var def)) return;
            byte slot = def.EquipSlot;
            if (slot >= EquipSlotId.Count) return; // not equippable (255 or out of range)

            var bag = state.EntityManager.GetBuffer<InventorySlot>(player);
            if (InventoryMath.CountOf(bag, itemId) < 1) return; // the sender isn't carrying it

            var slots = state.EntityManager.GetBuffer<EquipmentSlot>(player);
            ushort oldItem = slots[slot].ItemId;

            // Atomicity: if the slot is occupied, the bag MUST be able to hold the swapped-out item before we
            // touch anything; reject the whole equip otherwise so the old item is never lost.
            if (oldItem != 0 && !InventoryMath.CanDeposit(bag, oldItem, 1, StackMaxOf(ref db, oldItem), Tuning.InventoryMaxSlots))
                return;

            // Commit in-place.
            InventoryMath.Withdraw(bag, itemId, 1);
            if (oldItem != 0)
            {
                InventoryMath.Deposit(bag, oldItem, 1, StackMaxOf(ref db, oldItem), Tuning.InventoryMaxSlots);
                StripSlotEffects(ref state, player, slot);
            }

            slots[slot] = new EquipmentSlot { ItemId = itemId };
            ApplySlotEffects(ref state, player, slot, def);
        }

        static void HandleUnequip(ref SystemState state, ItemDatabase itemDb, Entity player, byte slot)
        {
            if (slot >= EquipSlotId.Count) return;
            ref var db = ref itemDb.Value.Value;

            var slots = state.EntityManager.GetBuffer<EquipmentSlot>(player);
            ushort item = slots[slot].ItemId;
            if (item == 0) return; // nothing equipped

            var bag = state.EntityManager.GetBuffer<InventorySlot>(player);
            if (!InventoryMath.CanDeposit(bag, item, 1, StackMaxOf(ref db, item), Tuning.InventoryMaxSlots))
                return; // bag full -> can't unequip (no item loss)

            InventoryMath.Deposit(bag, item, 1, StackMaxOf(ref db, item), Tuning.InventoryMaxSlots);
            slots[slot] = new EquipmentSlot { ItemId = 0 };
            StripSlotEffects(ref state, player, slot);
        }

        static void ApplySlotEffects(ref SystemState state, Entity player, byte slot, ItemDefBlob def)
        {
            // Weapon slot drives the active ability (swaps prefab + base stats via StatRecomputeSystem).
            if (slot == EquipSlotId.Weapon && def.GrantedAbilityId != 0)
                state.EntityManager.SetComponentData(player, new AbilityRef { Id = def.GrantedAbilityId });

            var mods = state.EntityManager.GetBuffer<StatModifier>(player);
            uint sourceId = Tuning.EquipSourceIdBase + (uint)slot;
            for (int i = 0; i < ItemDefBlob.MaxMods; i++)
            {
                var m = def.GetMod(i);
                if (m.Target == 255) continue;
                mods.Add(new StatModifier { Target = m.Target, Op = m.Op, Value = m.Value, SourceId = sourceId });
            }
        }

        static void StripSlotEffects(ref SystemState state, Entity player, byte slot)
        {
            var mods = state.EntityManager.GetBuffer<StatModifier>(player);
            TimedModifierUtil.RemoveBySourceId(mods, Tuning.EquipSourceIdBase + (uint)slot);

            // Weapon slot: restore the unarmed/base ability.
            if (slot == EquipSlotId.Weapon)
            {
                byte fallback = state.EntityManager.GetComponentData<DefaultAbility>(player).Id;
                state.EntityManager.SetComponentData(player, new AbilityRef { Id = fallback });
            }
        }

        static int StackMaxOf(ref ItemDatabaseBlob db, ushort itemId)
            => db.TryGetItem(itemId, out var d) && d.StackMax > 0 ? d.StackMax : 1;
    }
}