tbd-station-14/Content.Shared/MobState/Components/MobStateComponent.cs

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using Content.Shared.Alert;
using Content.Shared.Damage;
using Content.Shared.FixedPoint;
using Content.Shared.MobState.State;
using Robust.Shared.GameObjects;
using Robust.Shared.GameStates;
using Robust.Shared.IoC;
using Robust.Shared.Players;
using Robust.Shared.Serialization;
using Robust.Shared.Serialization.Manager.Attributes;
using Robust.Shared.ViewVariables;

namespace Content.Shared.MobState.Components
{
    /// <summary>
    ///     When attached to an <see cref="DamageableComponent"/>,
    ///     this component will handle critical and death behaviors for mobs.
    ///     Additionally, it handles sending effects to clients
    ///     (such as blur effect for unconsciousness) and managing the health HUD.
    /// </summary>
    [RegisterComponent]
    [NetworkedComponent]
    public class MobStateComponent : Component
    {
        public override string Name => "MobState";

        /// <summary>
        ///     States that this <see cref="MobStateComponent"/> mapped to
        ///     the amount of damage at which they are triggered.
        ///     A threshold is reached when the total damage of an entity is equal
        ///     to or higher than the int key, but lower than the next threshold.
        ///     Ordered from lowest to highest.
        /// </summary>
        [ViewVariables]
        [DataField("thresholds")]
        private readonly SortedDictionary<int, IMobState> _lowestToHighestStates = new();

        // TODO Remove Nullability?
        [ViewVariables]
        public IMobState? CurrentState { get; private set; }

        [ViewVariables]
        public FixedPoint2? CurrentThreshold { get; private set; }

        public IEnumerable<KeyValuePair<int, IMobState>> _highestToLowestStates => _lowestToHighestStates.Reverse();

        protected override void Startup()
        {
            base.Startup();

            if (CurrentState != null && CurrentThreshold != null)
            {
                // Initialize with given states
                SetMobState(null, (CurrentState, CurrentThreshold.Value));
            }
            else
            {
                // Initialize with some amount of damage, defaulting to 0.
                UpdateState(Owner.GetComponentOrNull<DamageableComponent>()?.TotalDamage ?? FixedPoint2.Zero);
            }
        }

        protected override void OnRemove()
        {
            if (IoCManager.Resolve<IEntityManager>().TryGetComponent(Owner.Uid, out SharedAlertsComponent? status))
            {
                status.ClearAlert(AlertType.HumanHealth);
            }

            base.OnRemove();
        }

        public override ComponentState GetComponentState()
        {
            return new MobStateComponentState(CurrentThreshold);
        }

        public override void HandleComponentState(ComponentState? curState, ComponentState? nextState)
        {
            base.HandleComponentState(curState, nextState);

            if (curState is not MobStateComponentState state)
            {
                return;
            }

            if (CurrentThreshold == state.CurrentThreshold)
            {
                return;
            }

            if (state.CurrentThreshold == null)
            {
                RemoveState();
            }
            else
            {
                UpdateState(state.CurrentThreshold.Value);
            }
        }

        public bool IsAlive()
        {
            return CurrentState?.IsAlive() ?? false;
        }

        public bool IsCritical()
        {
            return CurrentState?.IsCritical() ?? false;
        }

        public bool IsDead()
        {
            return CurrentState?.IsDead() ?? false;
        }

        public bool IsIncapacitated()
        {
            return CurrentState?.IsIncapacitated() ?? false;
        }

        public (IMobState state, FixedPoint2 threshold)? GetState(FixedPoint2 damage)
        {
            foreach (var (threshold, state) in _highestToLowestStates)
            {
                if (damage >= threshold)
                {
                    return (state, threshold);
                }
            }

            return null;
        }

        public bool TryGetState(
            FixedPoint2 damage,
            [NotNullWhen(true)] out IMobState? state,
            out FixedPoint2 threshold)
        {
            var highestState = GetState(damage);

            if (highestState == null)
            {
                state = default;
                threshold = default;
                return false;
            }

            (state, threshold) = highestState.Value;
            return true;
        }

        private (IMobState state, FixedPoint2 threshold)? GetEarliestState(FixedPoint2 minimumDamage, Predicate<IMobState> predicate)
        {
            foreach (var (threshold, state) in _lowestToHighestStates)
            {
                if (threshold < minimumDamage ||
                    !predicate(state))
                {
                    continue;
                }

                return (state, threshold);
            }

            return null;
        }

        private (IMobState state, FixedPoint2 threshold)? GetPreviousState(FixedPoint2 maximumDamage, Predicate<IMobState> predicate)
        {
            foreach (var (threshold, state) in _highestToLowestStates)
            {
                if (threshold > maximumDamage ||
                    !predicate(state))
                {
                    continue;
                }

                return (state, threshold);
            }

            return null;
        }

        public (IMobState state, FixedPoint2 threshold)? GetEarliestCriticalState(FixedPoint2 minimumDamage)
        {
            return GetEarliestState(minimumDamage, s => s.IsCritical());
        }

        public (IMobState state, FixedPoint2 threshold)? GetEarliestIncapacitatedState(FixedPoint2 minimumDamage)
        {
            return GetEarliestState(minimumDamage, s => s.IsIncapacitated());
        }

        public (IMobState state, FixedPoint2 threshold)? GetEarliestDeadState(FixedPoint2 minimumDamage)
        {
            return GetEarliestState(minimumDamage, s => s.IsDead());
        }

        public (IMobState state, FixedPoint2 threshold)? GetPreviousCriticalState(FixedPoint2 minimumDamage)
        {
            return GetPreviousState(minimumDamage, s => s.IsCritical());
        }

        private bool TryGetState(
            (IMobState state, FixedPoint2 threshold)? tuple,
            [NotNullWhen(true)] out IMobState? state,
            out FixedPoint2 threshold)
        {
            if (tuple == null)
            {
                state = default;
                threshold = default;
                return false;
            }

            (state, threshold) = tuple.Value;
            return true;
        }

        public bool TryGetEarliestCriticalState(
            FixedPoint2 minimumDamage,
            [NotNullWhen(true)] out IMobState? state,
            out FixedPoint2 threshold)
        {
            var earliestState = GetEarliestCriticalState(minimumDamage);

            return TryGetState(earliestState, out state, out threshold);
        }

        public bool TryGetEarliestIncapacitatedState(
            FixedPoint2 minimumDamage,
            [NotNullWhen(true)] out IMobState? state,
            out FixedPoint2 threshold)
        {
            var earliestState = GetEarliestIncapacitatedState(minimumDamage);

            return TryGetState(earliestState, out state, out threshold);
        }

        public bool TryGetEarliestDeadState(
            FixedPoint2 minimumDamage,
            [NotNullWhen(true)] out IMobState? state,
            out FixedPoint2 threshold)
        {
            var earliestState = GetEarliestDeadState(minimumDamage);

            return TryGetState(earliestState, out state, out threshold);
        }

        public bool TryGetPreviousCriticalState(
            FixedPoint2 maximumDamage,
            [NotNullWhen(true)] out IMobState? state,
            out FixedPoint2 threshold)
        {
            var earliestState = GetPreviousCriticalState(maximumDamage);

            return TryGetState(earliestState, out state, out threshold);
        }

        private void RemoveState()
        {
            var old = CurrentState;
            CurrentState = null;
            CurrentThreshold = null;

            SetMobState(old, null);
        }

        /// <summary>
        ///     Updates the mob state..
        /// </summary>
        public void UpdateState(FixedPoint2 damage)
        {
            if (!TryGetState(damage, out var newState, out var threshold))
            {
                return;
            }

            SetMobState(CurrentState, (newState, threshold));
        }

        /// <summary>
        ///     Sets the mob state and marks the component as dirty.
        /// </summary>
        private void SetMobState(IMobState? old, (IMobState state, FixedPoint2 threshold)? current)
        {
            if (!current.HasValue)
            {
                old?.ExitState(OwnerUid, IoCManager.Resolve<IEntityManager>());
                return;
            }

            var (state, threshold) = current.Value;

            CurrentThreshold = threshold;

            if (state == old)
            {
                state.UpdateState(OwnerUid, threshold, IoCManager.Resolve<IEntityManager>());
                return;
            }

            old?.ExitState(OwnerUid, IoCManager.Resolve<IEntityManager>());

            CurrentState = state;

            state.EnterState(OwnerUid, IoCManager.Resolve<IEntityManager>());
            state.UpdateState(OwnerUid, threshold, IoCManager.Resolve<IEntityManager>());

            var message = new MobStateChangedMessage(this, old, state);
#pragma warning disable 618
            SendMessage(message);
#pragma warning restore 618
            IoCManager.Resolve<IEntityManager>().EventBus.RaiseEvent(EventSource.Local, message);

            Dirty();
        }
    }

    [Serializable, NetSerializable]
    public class MobStateComponentState : ComponentState
    {
        public readonly FixedPoint2? CurrentThreshold;

        public MobStateComponentState(FixedPoint2? currentThreshold)
        {
            CurrentThreshold = currentThreshold;
        }
    }
}