BeatSaber/Assets/VRBeatsKit/Scripts/Core/SaberTrailEffect.cs

using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Rendering;

namespace VRBeats
{
    /// <summary>
    /// Builds a short world-space ribbon from the whole saber blade instead of
    /// attaching a TrailRenderer to the tip. This keeps the afterimage on the
    /// blade surface and avoids the "propeller" look from a single end point.
    /// </summary>
    public class SaberTrailEffect : MonoBehaviour
    {
        private const int MaxSamples = 18;
        private const float TrailLifetime = 0.13f;
        private const float MinSwingSpeed = 0.75f;
        private const float MinSampleDistance = 0.012f;
        private const float BladeBaseT = 0.10f;
        private const float BladeTipT = 0.98f;

        private readonly List<BladeSample> samples = new List<BladeSample>(MaxSamples);
        private readonly List<Vector3> vertices = new List<Vector3>(MaxSamples * 2);
        private readonly List<Color> colors = new List<Color>(MaxSamples * 2);
        private readonly List<Vector2> uvs = new List<Vector2>(MaxSamples * 2);
        private readonly List<int> triangles = new List<int>((MaxSamples - 1) * 12);

        private Transform bladeStart;
        private Transform bladeEnd;
        private GameObject wideObject;
        private GameObject coreObject;
        private Mesh wideMesh;
        private Mesh coreMesh;
        private MeshRenderer wideRenderer;
        private MeshRenderer coreRenderer;
        private Material wideMaterial;
        private Material coreMaterial;

        private Color trailColor = Color.cyan;
        private bool visible = true;
        private bool hasLastFrame;
        private Vector3 lastBase;
        private Vector3 lastTip;

        private struct BladeSample
        {
            public Vector3 basePos;
            public Vector3 tipPos;
            public float time;
        }

        private void Awake()
        {
            RemoveLegacyTrailRenderers();
            ResolveBladeAnchors();
            EnsureRenderers();
            SetColor(trailColor);
            SetVisible(visible);
        }

        private void LateUpdate()
        {
            if (!visible)
            {
                Clear();
                hasLastFrame = false;
                return;
            }

            if (bladeEnd == null)
            {
                ResolveBladeAnchors();
            }

            Vector3 basePos;
            Vector3 tipPos;
            GetBladeSegment(out basePos, out tipPos);

            float now = Time.time;
            float speed = 0f;
            float moved = 0f;

            if (hasLastFrame)
            {
                float deltaTime = Mathf.Max(Time.deltaTime, 0.0001f);
                speed = Mathf.Max(
                    Vector3.Distance(basePos, lastBase),
                    Vector3.Distance(tipPos, lastTip)) / deltaTime;
                moved = Mathf.Max(
                    Vector3.Distance(basePos, lastBase),
                    Vector3.Distance(tipPos, lastTip));
            }

            if (!hasLastFrame || (speed >= MinSwingSpeed && moved >= MinSampleDistance))
            {
                AddSample(basePos, tipPos, now);
            }

            TrimExpiredSamples(now);
            BuildRibbon(wideMesh, now, 0f, 1f, trailColor, 0.10f, 0.36f);

            Color coreColor = Color.Lerp(Color.white, trailColor, CoreColorWeight(trailColor));
            float alphaMultiplier = VisibilityAlphaMultiplier(trailColor);
            BuildRibbon(coreMesh, now, 0.42f, 1f, coreColor, 0.14f * alphaMultiplier, 0.48f * alphaMultiplier);

            lastBase = basePos;
            lastTip = tipPos;
            hasLastFrame = true;
        }

        private void OnDisable()
        {
            Clear();
            hasLastFrame = false;
        }

        private void OnDestroy()
        {
            DestroyRuntimeObject(wideObject);
            DestroyRuntimeObject(coreObject);
            DestroyRuntimeObject(wideMesh);
            DestroyRuntimeObject(coreMesh);
            DestroyRuntimeObject(wideMaterial);
            DestroyRuntimeObject(coreMaterial);
        }

        public void SetVisible(bool value)
        {
            visible = value;

            if (wideRenderer != null)
            {
                wideRenderer.enabled = value;
            }

            if (coreRenderer != null)
            {
                coreRenderer.enabled = value;
            }

            if (!value)
            {
                Clear();
                hasLastFrame = false;
            }
        }

        public void SetColor(Color color)
        {
            trailColor = NormalizeColor(color);
            EnsureRenderers();
            float alphaMultiplier = VisibilityAlphaMultiplier(trailColor);
            ApplyMaterialColor(wideMaterial, trailColor, 0.34f * alphaMultiplier);

            Color coreColor = Color.Lerp(Color.white, trailColor, CoreColorWeight(trailColor));
            ApplyMaterialColor(coreMaterial, coreColor, 0.50f * alphaMultiplier);
        }

        private void ResolveBladeAnchors()
        {
            bladeStart = FindChildRecursive(transform, "Start");
            bladeEnd = FindChildRecursive(transform, "End");

            if (bladeEnd == null)
            {
                bladeEnd = transform;
            }

            if (bladeStart == null)
            {
                bladeStart = transform;
            }
        }

        private void GetBladeSegment(out Vector3 basePos, out Vector3 tipPos)
        {
            Vector3 start = bladeStart != null ? bladeStart.position : transform.position;
            Vector3 end = bladeEnd != null ? bladeEnd.position : transform.position + transform.forward;

            if (Vector3.Distance(start, end) < 0.001f)
            {
                end = start + transform.forward * 0.8f;
            }

            basePos = Vector3.Lerp(start, end, BladeBaseT);
            tipPos = Vector3.Lerp(start, end, BladeTipT);
        }

        private void AddSample(Vector3 basePos, Vector3 tipPos, float time)
        {
            if (samples.Count >= MaxSamples)
            {
                samples.RemoveAt(0);
            }

            samples.Add(new BladeSample
            {
                basePos = basePos,
                tipPos = tipPos,
                time = time
            });
        }

        private void TrimExpiredSamples(float now)
        {
            for (int i = samples.Count - 1; i >= 0; i--)
            {
                if (now - samples[i].time > TrailLifetime)
                {
                    samples.RemoveAt(i);
                }
            }
        }

        private void BuildRibbon(Mesh mesh, float now, float innerT, float outerT, Color color, float innerAlpha, float outerAlpha)
        {
            if (mesh == null)
            {
                return;
            }

            vertices.Clear();
            colors.Clear();
            uvs.Clear();
            triangles.Clear();

            if (samples.Count < 2)
            {
                mesh.Clear();
                return;
            }

            for (int i = 0; i < samples.Count; i++)
            {
                BladeSample sample = samples[i];
                float age01 = Mathf.Clamp01((now - sample.time) / TrailLifetime);
                float fade = Mathf.Pow(1f - age01, 1.7f);
                float along = samples.Count <= 1 ? 0f : i / (float)(samples.Count - 1);

                Vector3 inner = Vector3.Lerp(sample.basePos, sample.tipPos, innerT);
                Vector3 outer = Vector3.Lerp(sample.basePos, sample.tipPos, outerT);

                vertices.Add(inner);
                vertices.Add(outer);
                colors.Add(WithAlpha(color, fade * innerAlpha));
                colors.Add(WithAlpha(color, fade * outerAlpha));
                uvs.Add(new Vector2(0f, along));
                uvs.Add(new Vector2(1f, along));
            }

            for (int i = 0; i < samples.Count - 1; i++)
            {
                int a = i * 2;
                int b = a + 1;
                int c = a + 2;
                int d = a + 3;

                triangles.Add(a);
                triangles.Add(b);
                triangles.Add(c);
                triangles.Add(b);
                triangles.Add(d);
                triangles.Add(c);

                triangles.Add(c);
                triangles.Add(b);
                triangles.Add(a);
                triangles.Add(c);
                triangles.Add(d);
                triangles.Add(b);
            }

            mesh.Clear();
            mesh.SetVertices(vertices);
            mesh.SetColors(colors);
            mesh.SetUVs(0, uvs);
            mesh.SetTriangles(triangles, 0);
            mesh.RecalculateBounds();
        }

        private void EnsureRenderers()
        {
            if (wideObject == null)
            {
                wideRenderer = CreateRibbonObject("SaberBladeAfterimage_Wide", out wideObject, out wideMesh, out wideMaterial, false);
            }

            if (coreObject == null)
            {
                coreRenderer = CreateRibbonObject("SaberBladeAfterimage_Core", out coreObject, out coreMesh, out coreMaterial, true);
            }
        }

        private MeshRenderer CreateRibbonObject(string objectName, out GameObject ribbonObject, out Mesh mesh, out Material material, bool additive)
        {
            ribbonObject = new GameObject(objectName);
            ribbonObject.hideFlags = HideFlags.DontSave;
            ribbonObject.transform.SetPositionAndRotation(Vector3.zero, Quaternion.identity);
            ribbonObject.transform.localScale = Vector3.one;

            mesh = new Mesh
            {
                name = objectName + " Mesh"
            };
            mesh.MarkDynamic();

            MeshFilter filter = ribbonObject.AddComponent<MeshFilter>();
            filter.sharedMesh = mesh;

            MeshRenderer renderer = ribbonObject.AddComponent<MeshRenderer>();
            material = CreateTrailMaterial(objectName + " Material", additive);
            renderer.sharedMaterial = material;
            renderer.shadowCastingMode = ShadowCastingMode.Off;
            renderer.receiveShadows = false;
            renderer.motionVectorGenerationMode = MotionVectorGenerationMode.ForceNoMotion;
            renderer.enabled = visible;

            return renderer;
        }

        private Material CreateTrailMaterial(string materialName, bool additive)
        {
            Shader shader = Shader.Find("Sprites/Default");
            if (shader == null)
            {
                shader = Shader.Find("Unlit/Transparent");
            }

            Material material = new Material(shader)
            {
                name = materialName,
                hideFlags = HideFlags.DontSave,
                renderQueue = (int)RenderQueue.Transparent
            };

            if (material.HasProperty("_SrcBlend"))
            {
                material.SetInt("_SrcBlend", additive ? (int)BlendMode.SrcAlpha : (int)BlendMode.SrcAlpha);
            }

            if (material.HasProperty("_DstBlend"))
            {
                material.SetInt("_DstBlend", additive ? (int)BlendMode.One : (int)BlendMode.OneMinusSrcAlpha);
            }

            if (material.HasProperty("_ZWrite"))
            {
                material.SetInt("_ZWrite", 0);
            }

            return material;
        }

        private void ApplyMaterialColor(Material material, Color color, float alpha)
        {
            if (material == null)
            {
                return;
            }

            Color materialColor = WithAlpha(color, alpha);
            if (material.HasProperty("_Color"))
            {
                material.SetColor("_Color", materialColor);
            }
        }

        private void Clear()
        {
            samples.Clear();

            if (wideMesh != null)
            {
                wideMesh.Clear();
            }

            if (coreMesh != null)
            {
                coreMesh.Clear();
            }
        }

        private void RemoveLegacyTrailRenderers()
        {
            TrailRenderer[] legacyTrails = GetComponentsInChildren<TrailRenderer>(true);
            for (int i = legacyTrails.Length - 1; i >= 0; i--)
            {
                DestroyRuntimeObject(legacyTrails[i]);
            }
        }

        private static Transform FindChildRecursive(Transform root, string childName)
        {
            if (root == null)
            {
                return null;
            }

            if (root.name == childName)
            {
                return root;
            }

            for (int i = 0; i < root.childCount; i++)
            {
                Transform found = FindChildRecursive(root.GetChild(i), childName);
                if (found != null)
                {
                    return found;
                }
            }

            return null;
        }

        private static Color NormalizeColor(Color color)
        {
            float max = Mathf.Max(color.r, color.g, color.b);
            if (max > 1f)
            {
                color.r /= max;
                color.g /= max;
                color.b /= max;
            }

            color.a = 1f;
            return color;
        }

        private static float CoreColorWeight(Color color)
        {
            return IsBlueDominant(color) ? 0.78f : 0.45f;
        }

        private static float VisibilityAlphaMultiplier(Color color)
        {
            return IsBlueDominant(color) ? 1.35f : 1f;
        }

        private static bool IsBlueDominant(Color color)
        {
            return color.b > color.r && color.b >= color.g;
        }

        private static Color WithAlpha(Color color, float alpha)
        {
            color.a = Mathf.Clamp01(alpha);
            return color;
        }

        private static void DestroyRuntimeObject(Object target)
        {
            if (target == null)
            {
                return;
            }

            if (Application.isPlaying)
            {
                Destroy(target);
            }
            else
            {
                DestroyImmediate(target);
            }
        }
    }
}