A-Fairytale-Gone-Bad-Briar-Queen/Assets/Scripts/Framework/Effects/UIDissolveImage.cs

using System;
using System.Collections.Generic;
using System.Threading;
using BriarQueen.Framework.Services.Destruction;
using Cysharp.Threading.Tasks;
using PrimeTween;
using TMPro;
using UnityEngine;
using UnityEngine.UI;
using VContainer;

namespace BriarQueen.Framework.Effects
{
    [ExecuteAlways]
    public class UIDissolveImage : MonoBehaviour
    {
        private static readonly int _dissolveAmountId = Shader.PropertyToID("_DissolveAmount");
        private static readonly int _reverseDirectionId = Shader.PropertyToID("_ReverseDirection");

        [Header("References")]
        [SerializeField]
        private Image _image;

        [SerializeField]
        private CanvasGroup _canvasGroup;

        [SerializeField]
        private Material _dissolveMaterialTemplate;

        [Header("Targeting")]
        [SerializeField]
        private bool _dissolveChildGraphics;

        [SerializeField]
        private bool _includeInactiveChildGraphics = true;

        [SerializeField]
        private bool _includeTextGraphics;

        [Header("Tween")]
        [SerializeField]
        private float _duration = 0.75f;

        [SerializeField]
        private Ease _ease = Ease.InOutSine;

        [SerializeField]
        private bool _useUnscaledTime = true;

        [Header("Direction")]
        [SerializeField]
        private bool _reverseDirection;

        [Header("Destruction")]
        [SerializeField]
        private bool _destroyWhenFullyDissolved = true;

        [Header("Editor Preview")]
        [SerializeField]
        private bool _previewInEditMode;

        [SerializeField]
        [Range(0f, 1f)]
        private float _previewDissolveAmount;

        private readonly List<Graphic> _targetGraphics = new();
        private readonly List<Material> _originalMaterials = new();

        private CancellationTokenSource _dissolveCts;
        private DestructionService _destructionService;
        private bool _hasOriginalMaterials;
        private bool _isPreviewMaterial;
        private Material _runtimeMaterial;
        private Sequence _dissolveSequence;

        public float DissolveAmount
        {
            get
            {
                if (_runtimeMaterial == null)
                {
                    return 0f;
                }

                return _runtimeMaterial.GetFloat(_dissolveAmountId);
            }
            set => SetDissolveAmount(value);
        }

        [Inject]
        public void Construct(DestructionService destructionService)
        {
            _destructionService = destructionService;
        }

        private void Awake()
        {
            ResolveReferences();

            if (Application.isPlaying)
            {
                CreateRuntimeMaterial(false);
                SetDissolveAmount(0f);
                return;
            }

#if UNITY_EDITOR
            RefreshEditModePreview();
#endif
        }

        private void OnEnable()
        {
#if UNITY_EDITOR
            if (!Application.isPlaying)
            {
                RefreshEditModePreview();
            }
#endif
        }

        private void OnDisable()
        {
#if UNITY_EDITOR
            if (!Application.isPlaying)
            {
                RestoreEditModeMaterial();
            }
#endif
        }

        private void OnDestroy()
        {
            CancelDissolve();

            if (Application.isPlaying)
            {
                DestroyRuntimeMaterial(false);
                return;
            }

#if UNITY_EDITOR
            RestoreEditModeMaterial();
#endif
        }

        private void OnValidate()
        {
            ResolveReferences();

            if (Application.isPlaying)
            {
                SetReverseDirection(_reverseDirection);
                return;
            }

#if UNITY_EDITOR
            RestoreEditModeMaterial();
            RefreshEditModePreview();
#endif
        }

        public UniTask DissolveIn()
        {
            gameObject.SetActive(true);
            return TweenDissolve(1f, 0f, _duration, false);
        }

        public UniTask DissolveOut()
        {
            return TweenDissolve(0f, 1f, _duration, _destroyWhenFullyDissolved);
        }

        public UniTask DissolveOut(bool destroyWhenComplete)
        {
            return TweenDissolve(0f, 1f, _duration, destroyWhenComplete);
        }

        public UniTask DissolveOutAndDestroy()
        {
            return TweenDissolve(0f, 1f, _duration, true);
        }

        public UniTask TweenDissolve(float from, float to, float duration)
        {
            return TweenDissolve(from, to, duration, false);
        }

        public async UniTask TweenDissolve(float from, float to, float duration, bool destroyWhenComplete)
        {
            if (_runtimeMaterial == null)
            {
                CreateRuntimeMaterial(false);
            }

            CancelDissolve();

            _dissolveCts = new CancellationTokenSource();

            SetDissolveAmount(from);

            _dissolveSequence = Sequence.Create(useUnscaledTime: _useUnscaledTime)
                .Group(Tween.Custom(
                    from,
                    to,
                    Mathf.Max(0f, duration),
                    SetDissolveAmount,
                    _ease,
                    useUnscaledTime: _useUnscaledTime));

            try
            {
                await _dissolveSequence.ToUniTask(cancellationToken: _dissolveCts.Token);
                SetDissolveAmount(to);

                if (destroyWhenComplete && Mathf.Approximately(to, 1f))
                {
                    await DestroyAfterDissolve();
                }
            }
            catch (OperationCanceledException)
            {
                // Interrupted by another dissolve request or object destruction.
            }
            finally
            {
                _dissolveSequence = default;

                if (_dissolveCts != null)
                {
                    _dissolveCts.Dispose();
                    _dissolveCts = null;
                }
            }
        }

        public void SetDissolveAmount(float amount)
        {
            if (_runtimeMaterial == null)
            {
                return;
            }

            _runtimeMaterial.SetFloat(_dissolveAmountId, Mathf.Clamp01(amount));
            SetTargetsMaterialDirty();
        }

        public void SetReverseDirection(bool reverseDirection)
        {
            _reverseDirection = reverseDirection;

            if (_runtimeMaterial == null)
            {
                return;
            }

            _runtimeMaterial.SetFloat(_reverseDirectionId, _reverseDirection ? 1f : 0f);
            SetTargetsMaterialDirty();
        }

        public void CancelDissolve()
        {
            if (_dissolveSequence.isAlive)
            {
                _dissolveSequence.Stop();
                _dissolveSequence = default;
            }

            if (_dissolveCts != null)
            {
                _dissolveCts.Cancel();
                _dissolveCts.Dispose();
                _dissolveCts = null;
            }
        }

        private async UniTask DestroyAfterDissolve()
        {
            if (_destructionService != null)
            {
                await _destructionService.Destroy(gameObject);
                return;
            }

            Debug.LogWarning($"[{nameof(UIDissolveImage)}] Missing {nameof(DestructionService)}. Destroying directly.");
            Destroy(gameObject);
        }

        private void ResolveReferences()
        {
            if (_image == null)
            {
                _image = GetComponent<Image>();
            }

            if (_canvasGroup == null)
            {
                _canvasGroup = GetComponent<CanvasGroup>();
            }
        }

        private void ResolveTargetGraphics()
        {
            ResolveReferences();
            _targetGraphics.Clear();

            if (_dissolveChildGraphics)
            {
                var root = _canvasGroup != null ? _canvasGroup.transform : transform;
                var graphics = root.GetComponentsInChildren<Graphic>(_includeInactiveChildGraphics);

                foreach (var graphic in graphics)
                {
                    if (IsValidTargetGraphic(graphic))
                    {
                        _targetGraphics.Add(graphic);
                    }
                }

                return;
            }

            if (_image != null)
            {
                _targetGraphics.Add(_image);
            }
        }

        private bool IsValidTargetGraphic(Graphic graphic)
        {
            if (graphic == null)
            {
                return false;
            }

            // TMP needs a TMP-compatible dissolve shader; replacing its SDF material breaks text rendering.
            if (!_includeTextGraphics && graphic is TMP_Text)
            {
                return false;
            }

            return true;
        }

        private void CreateRuntimeMaterial(bool isPreviewMaterial)
        {
            ResolveTargetGraphics();

            if (_targetGraphics.Count == 0)
            {
                return;
            }

            var sourceMaterial = _dissolveMaterialTemplate != null
                ? _dissolveMaterialTemplate
                : _targetGraphics[0].material;

            if (sourceMaterial == null)
            {
                return;
            }

            if (_runtimeMaterial != null)
            {
                ApplyRuntimeMaterialToTargets();
                return;
            }

            SaveOriginalMaterials();

            _runtimeMaterial = Instantiate(sourceMaterial);
            _runtimeMaterial.name = isPreviewMaterial
                ? $"{nameof(UIDissolveImage)} Preview Material"
                : $"{nameof(UIDissolveImage)} Runtime Material";
            _isPreviewMaterial = isPreviewMaterial;

            if (isPreviewMaterial)
            {
                _runtimeMaterial.hideFlags = HideFlags.DontSaveInEditor;
            }

            ApplyRuntimeMaterialToTargets();
            SetReverseDirection(_reverseDirection);
        }

        private void ApplyRuntimeMaterialToTargets()
        {
            foreach (var graphic in _targetGraphics)
            {
                if (graphic == null)
                {
                    continue;
                }

                graphic.material = _runtimeMaterial;
                graphic.SetMaterialDirty();
            }
        }

        private void SaveOriginalMaterials()
        {
            _originalMaterials.Clear();

            foreach (var graphic in _targetGraphics)
            {
                _originalMaterials.Add(graphic != null ? graphic.material : null);
            }

            _hasOriginalMaterials = true;
        }

        private void RestoreOriginalMaterials()
        {
            if (!_hasOriginalMaterials)
            {
                return;
            }

            var count = Mathf.Min(_targetGraphics.Count, _originalMaterials.Count);
            for (var i = 0; i < count; i++)
            {
                var graphic = _targetGraphics[i];
                if (graphic == null)
                {
                    continue;
                }

                graphic.material = _originalMaterials[i];
                graphic.SetMaterialDirty();
            }

            _originalMaterials.Clear();
            _hasOriginalMaterials = false;
        }

        private void SetTargetsMaterialDirty()
        {
            foreach (var graphic in _targetGraphics)
            {
                if (graphic != null)
                {
                    graphic.SetMaterialDirty();
                }
            }
        }

        private void DestroyRuntimeMaterial(bool immediate)
        {
            if (_runtimeMaterial == null)
            {
                return;
            }

            RestoreOriginalMaterials();

            if (immediate)
            {
                DestroyImmediate(_runtimeMaterial);
            }
            else
            {
                Destroy(_runtimeMaterial);
            }

            _runtimeMaterial = null;
            _isPreviewMaterial = false;
        }

#if UNITY_EDITOR
        private void RefreshEditModePreview()
        {
            if (Application.isPlaying)
            {
                return;
            }

            if (!_previewInEditMode)
            {
                RestoreEditModeMaterial();
                return;
            }

            CreateRuntimeMaterial(true);
            SetReverseDirection(_reverseDirection);
            SetDissolveAmount(_previewDissolveAmount);
        }

        private void RestoreEditModeMaterial()
        {
            if (!_isPreviewMaterial)
            {
                return;
            }

            DestroyRuntimeMaterial(true);
        }
#endif
    }
}