Project-M/Assets/_Project/Scripts/Client/Presentation/EnemyHitFlashSystem.cs

using System.Collections.Generic;
using ProjectM.Simulation;
using Unity.Entities;
using Unity.Mathematics;
using Unity.Rendering;       // URPMaterialPropertyBaseColor, MaterialMeshInfo (Entities Graphics)
using UnityEngine;

namespace ProjectM.Client
{
    /// <summary>
    /// Client-only TRUE BODY hit-flash for enemies (the focused follow-up DR-041 deferred as "its own ShaderGraph
    /// slice"). Enemies render via Rukhanka GPU deformation: the ghost ROOT holds the gameplay components
    /// (<see cref="Health"/>, <see cref="EnemyTag"/>) while the visible meshes are LinkedEntityGroup CHILD render
    /// entities (each with a <see cref="MaterialMeshInfo"/> + the AnimatedLitShader material, whose <c>_BaseColor</c>
    /// is white). <see cref="CombatFeedbackSystem"/>'s colored particle puff is the asset-free stand-in; THIS system
    /// flashes the actual body by driving the built-in Entities-Graphics per-instance override
    /// <see cref="URPMaterialPropertyBaseColor"/> (a registered <c>[MaterialProperty("_BaseColor")]</c>) on those
    /// render children: on an enemy <see cref="Health"/>-decrease edge it lerps <c>_BaseColor</c> toward
    /// <see cref="FeelConfig.BodyFlashColor"/> and decays back to white. No new component type, NO ShaderGraph edit,
    /// no server work, no <c>[GhostField]</c> — observe-only client presentation, so it is rollback-irrelevant.
    /// The render children gain <see cref="URPMaterialPropertyBaseColor"/> lazily (added once per enemy via ECB);
    /// white is the baked rest value (every enemy uses the same AnimatedLitShader/Synty-atlas convention), so
    /// <c>Flash==0</c> restores the untouched look and the override is never visible at rest.
    /// </summary>
    [WorldSystemFilter(WorldSystemFilterFlags.ClientSimulation)]
    [UpdateInGroup(typeof(PresentationSystemGroup))]
    public partial class EnemyHitFlashSystem : SystemBase
    {
        class FlashEntry
        {
            public readonly List<Entity> RenderKids = new();
            public float LastHp;
            public float Flash;   // 1 on a fresh hit, decays to 0
            public bool Settled;  // wrote the final white frame after a flash ended (skips per-frame writes at rest)
            public bool IsPlayer; // C2: a player entry flashes PlayerHurtFlashColor (red), an enemy the white overdrive rest)
        }

        readonly Dictionary<Entity, FlashEntry> _tracked = new();
        readonly HashSet<Entity> _seen = new();
        readonly List<Entity> _stale = new();

        static readonly float4 White = new float4(1f, 1f, 1f, 1f);

        protected override void OnUpdate()
        {
            if (!FeelConfig.BodyFlashEnabled) { RestoreAllToRest(); return; } // settle any mid-flash body back to white before bailing
            float dt = SystemAPI.Time.DeltaTime;
            EntityManager.CompleteDependencyBeforeRO<Health>();

            // Pass 1: discover enemies, ensure each is tracked + its render children carry the override component.
            _seen.Clear();
            // Track B: created lazily — it only ever records on a newly-seen enemy, but it used to be built
            // and played back (a sync point) every single frame.
            EntityCommandBuffer ecb = default;
            bool hasEcb = false;
            foreach (var (health, entity) in
                     SystemAPI.Query<RefRO<Health>>().WithAny<EnemyTag, PlayerTag>().WithAll<LinkedEntityGroup>().WithEntityAccess())
            {
                _seen.Add(entity);
                if (_tracked.ContainsKey(entity)) continue;

                var entry = new FlashEntry { LastHp = health.ValueRO.Current, Flash = 0f, Settled = true, IsPlayer = EntityManager.HasComponent<PlayerTag>(entity) };
                var leg = EntityManager.GetBuffer<LinkedEntityGroup>(entity);
                for (int i = 0; i < leg.Length; i++)
                {
                    var c = leg[i].Value;
                    if (!EntityManager.Exists(c) || !EntityManager.HasComponent<MaterialMeshInfo>(c)) continue;
                    entry.RenderKids.Add(c);
                    if (!EntityManager.HasComponent<URPMaterialPropertyBaseColor>(c))
                    {
                        if (!hasEcb) { ecb = new EntityCommandBuffer(Unity.Collections.Allocator.Temp); hasEcb = true; }
                        ecb.AddComponent(c, new URPMaterialPropertyBaseColor { Value = White });
                    }
                }
                // Render children can lag ghost instantiation a frame; only finalize once we actually found them (else retry next frame).
                if (entry.RenderKids.Count > 0) _tracked[entity] = entry;
            }
            // Only pay the structural-change sync point on the frames that actually recorded something
            // (i.e. a newly-seen enemy) instead of every frame.
            if (hasEcb) { ecb.Playback(EntityManager); ecb.Dispose(); }

            // Pass 2: edge-detect Health, drive + decay the flash, write _BaseColor to the render children.
            var bc = FeelConfig.BodyFlashColor;
            float4 enemyPeak = new float4(bc.r, bc.g, bc.b, bc.a);
            var pc = FeelConfig.PlayerHurtFlashColor;
            float4 playerPeak = new float4(pc.r, pc.g, pc.b, pc.a);
            float decay = dt / math.max(0.01f, FeelConfig.BodyFlashDurationSec);
            foreach (var kv in _tracked)
            {
                var entity = kv.Key;
                var entry = kv.Value;
                if (!_seen.Contains(entity)) continue; // despawned -> pruned below

                float cur = EntityManager.GetComponentData<Health>(entity).Current;
                if (cur < entry.LastHp - 0.5f) { entry.Flash = 1f; entry.Settled = false; }
                entry.LastHp = cur;

                if (entry.Flash <= 0f)
                {
                    if (!entry.Settled) { WriteColor(entry, White); entry.Settled = true; } // settle to baked white once
                    continue;
                }

                entry.Flash = math.max(0f, entry.Flash - decay);
                float4 peak = entry.IsPlayer ? playerPeak : enemyPeak;
                WriteColor(entry, math.lerp(White, peak, entry.Flash));
            }

            // Prune despawned enemies (their render children die with them; just drop the managed entry).
            if (_tracked.Count != _seen.Count)
            {
                _stale.Clear();
                foreach (var kv in _tracked) if (!_seen.Contains(kv.Key)) _stale.Add(kv.Key);
                for (int i = 0; i < _stale.Count; i++) _tracked.Remove(_stale[i]);
            }
        }

        // Settle every tracked enemy's body back to its baked white rest color and stop tracking — invoked when
        // BodyFlashEnabled is toggled OFF mid-flash so no body is left frozen at an overdriven tint (re-tracked on re-enable).
        void RestoreAllToRest()
        {
            foreach (var kv in _tracked) WriteColor(kv.Value, White);
            _tracked.Clear();
        }

        void WriteColor(FlashEntry entry, float4 col)
        {
            for (int i = 0; i < entry.RenderKids.Count; i++)
            {
                var c = entry.RenderKids[i];
                if (EntityManager.Exists(c) && EntityManager.HasComponent<URPMaterialPropertyBaseColor>(c))
                    EntityManager.SetComponentData(c, new URPMaterialPropertyBaseColor { Value = col });
            }
        }
    }
}