tbd-station-14/Content.Client/GameObjects/Components/DoAfterComponent.cs

#nullable enable
using System;
using System.Collections.Generic;
using Content.Client.GameObjects.EntitySystems.DoAfter;
using Content.Shared.GameObjects.Components;
using Robust.Shared.GameObjects;
using Robust.Shared.Interfaces.Network;
using Robust.Shared.Interfaces.Timing;
using Robust.Shared.IoC;
using Robust.Shared.Players;

namespace Content.Client.GameObjects.Components
{
    [RegisterComponent]
    public sealed class DoAfterComponent : SharedDoAfterComponent
    {
        public override string Name => "DoAfter";

        public IReadOnlyDictionary<byte, ClientDoAfter> DoAfters => _doAfters;
        private readonly Dictionary<byte, ClientDoAfter> _doAfters = new Dictionary<byte, ClientDoAfter>();
        
        public readonly List<(TimeSpan CancelTime, ClientDoAfter Message)> CancelledDoAfters = 
                     new List<(TimeSpan CancelTime, ClientDoAfter Message)>();

        public DoAfterGui? Gui { get; set; }

        public override void HandleNetworkMessage(ComponentMessage message, INetChannel netChannel, ICommonSession? session = null)
        {
            base.HandleNetworkMessage(message, netChannel, session);
            switch (message)
            {
                case CancelledDoAfterMessage msg:
                    Cancel(msg.ID);
                    break;
            }
        }

        public override void OnAdd()
        {
            base.OnAdd();
            Enable();
        }

        public override void OnRemove()
        {
            base.OnRemove();
            Disable();
        }

        /// <summary>
        ///     For handling PVS so we dispose of controls if they go out of range
        /// </summary>
        public void Enable()
        {
            if (Gui != null && !Gui.Disposed)
                return;
            
            Gui = new DoAfterGui {AttachedEntity = Owner};
            
            foreach (var (_, doAfter) in _doAfters)
            {
                Gui.AddDoAfter(doAfter);
            }

            foreach (var (_, cancelled) in CancelledDoAfters)
            {
                Gui.CancelDoAfter(cancelled.ID);
            }
        }

        public void Disable()
        {
            Gui?.Dispose();
        }

        public override void HandleComponentState(ComponentState? curState, ComponentState? nextState)
        {
            base.HandleComponentState(curState, nextState);
            if (!(curState is DoAfterComponentState state))
                return;
            
            _doAfters.Clear();

            foreach (var doAfter in state.DoAfters)
            {
                _doAfters.Add(doAfter.ID, doAfter);
            }
            
            if (Gui == null || Gui.Disposed)
                return;

            foreach (var (_, doAfter) in _doAfters)
            {
                Gui.AddDoAfter(doAfter);
            }
        }

        /// <summary>
        ///     Remove a DoAfter without showing a cancellation graphic.
        /// </summary>
        /// <param name="clientDoAfter"></param>
        public void Remove(ClientDoAfter clientDoAfter)
        {
            if (_doAfters.ContainsKey(clientDoAfter.ID))
                _doAfters.Remove(clientDoAfter.ID);

            var found = false;
            
            for (var i = CancelledDoAfters.Count - 1; i >= 0; i--)
            {
                var cancelled = CancelledDoAfters[i];

                if (cancelled.Message == clientDoAfter)
                {
                    CancelledDoAfters.RemoveAt(i);
                    found = true;
                    break;
                }
            }

            if (!found)
                _doAfters.Remove(clientDoAfter.ID);

            Gui?.RemoveDoAfter(clientDoAfter.ID);
        }

        /// <summary>
        ///     Mark a DoAfter as cancelled and show a cancellation graphic.
        /// </summary>
        ///     Actual removal is handled by DoAfterEntitySystem.
        /// <param name="id"></param>
        /// <param name="currentTime"></param>
        public void Cancel(byte id, TimeSpan? currentTime = null)
        {
            foreach (var (_, cancelled) in CancelledDoAfters)
            {
                if (cancelled.ID == id)
                    return;
            }

            if (!_doAfters.ContainsKey(id))
                return;
            
            var doAfterMessage = _doAfters[id];
            currentTime ??= IoCManager.Resolve<IGameTiming>().CurTime;
            CancelledDoAfters.Add((currentTime.Value, doAfterMessage));
            Gui?.CancelDoAfter(id);
        }
    }
}