tbd-station-14/Content.Client/Overlays/ShowHealthIconsSystem.cs

using Content.Shared.Atmos.Rotting;
using Content.Shared.Damage;
using Content.Shared.Inventory.Events;
using Content.Shared.Mobs.Components;
using Content.Shared.Overlays;
using Content.Shared.StatusIcon;
using Content.Shared.StatusIcon.Components;
using Robust.Shared.Prototypes;
using System.Linq;
using Content.Shared._Offbrand.Wounds; // Offbrand
using Content.Shared.Mobs; // Offbrand

namespace Content.Client.Overlays;

/// <summary>
/// Shows a healthy icon on mobs.
/// </summary>
public sealed class ShowHealthIconsSystem : EquipmentHudSystem<ShowHealthIconsComponent>
{
    [Dependency] private readonly IPrototypeManager _prototypeMan = default!;

    [ViewVariables]
    public HashSet<string> DamageContainers = new();

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

        SubscribeLocalEvent<DamageableComponent, GetStatusIconsEvent>(OnGetStatusIconsEvent);
        SubscribeLocalEvent<ShowHealthIconsComponent, AfterAutoHandleStateEvent>(OnHandleState);
    }

    protected override void UpdateInternal(RefreshEquipmentHudEvent<ShowHealthIconsComponent> component)
    {
        base.UpdateInternal(component);

        foreach (var damageContainerId in component.Components.SelectMany(x => x.DamageContainers))
        {
            DamageContainers.Add(damageContainerId);
        }
    }

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

        DamageContainers.Clear();
    }

    private void OnHandleState(Entity<ShowHealthIconsComponent> ent, ref AfterAutoHandleStateEvent args)
    {
        RefreshOverlay();
    }

    private void OnGetStatusIconsEvent(Entity<DamageableComponent> entity, ref GetStatusIconsEvent args)
    {
        if (!IsActive)
            return;

        var healthIcons = DecideHealthIcons(entity);

        args.StatusIcons.AddRange(healthIcons);
    }

    // Begin Offbrand
    private List<HealthIconPrototype> DecideBrainHealthIcons(Entity<BrainDamageComponent, BrainDamageThresholdsComponent> ent)
    {
        if (ent.Comp2.CurrentState == MobState.Dead)
        {
            return new() { _prototypeMan.Index(ent.Comp2.DeadIcon) };
        }

        var current = ent.Comp1.Damage;
        var max = ent.Comp1.MaxDamage;

        if (ent.Comp2.CurrentState == MobState.Critical || ent.Comp1.Oxygen == 0)
        {
            var amount = ent.Comp2.CriticalDamageIcons.Count;
            var idx = Math.Clamp((int)Math.Floor(amount - (amount / max.Double()) * current.Double()), 0, amount-1);
            return new() { _prototypeMan.Index(ent.Comp2.CriticalDamageIcons[idx]) };
        }
        else
        {
            var amount = ent.Comp2.AliveDamageIcons.Count;
            var idx = Math.Clamp((int)Math.Floor(amount - (amount / max.Double()) * current.Double()), 0, amount-1);
            return new() { _prototypeMan.Index(ent.Comp2.AliveDamageIcons[idx]) };
        }
    }
    // End Offbrand

    private IReadOnlyList<HealthIconPrototype> DecideHealthIcons(Entity<DamageableComponent> entity)
    {
        if (TryComp<BrainDamageComponent>(entity, out var brain) &&
            TryComp<BrainDamageThresholdsComponent>(entity, out var thresholds))
        {
            return DecideBrainHealthIcons((entity.Owner, brain, thresholds));
        }

        var damageableComponent = entity.Comp;

        if (damageableComponent.DamageContainerID == null ||
            !DamageContainers.Contains(damageableComponent.DamageContainerID))
        {
            return Array.Empty<HealthIconPrototype>();
        }

        var result = new List<HealthIconPrototype>();

        // Here you could check health status, diseases, mind status, etc. and pick a good icon, or multiple depending on whatever.
        if (damageableComponent?.DamageContainerID == "Biological")
        {
            if (TryComp<MobStateComponent>(entity, out var state))
            {
                // Since there is no MobState for a rotting mob, we have to deal with this case first.
                if (HasComp<RottingComponent>(entity) && _prototypeMan.Resolve(damageableComponent.RottingIcon, out var rottingIcon))
                    result.Add(rottingIcon);
                else if (damageableComponent.HealthIcons.TryGetValue(state.CurrentState, out var value) && _prototypeMan.Resolve(value, out var icon))
                    result.Add(icon);
            }
        }

        return result;
    }
}