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

namespace WheelTrail
{
    /// <summary>
    /// A world-space trail renderer that stores the contact surface normal per point and builds a stable ribbon
    /// from path tangent so vertex normals follow the surface without twist from the wheel transform.
    /// 
    /// made by StCost for COLLAPSE MACHINE. details in blog: https://dreamingsaints.github.io/blog/51-wheeltrailrenderer/
    /// </summary>
    [DisallowMultipleComponent]
    public class WheelTrailRenderer : MonoBehaviour
    {
        [SerializeField, Min(0.1f)] private float maxDistance = 50f;
        [SerializeField, Min(0.001f)] private float minVertexDistance = 0.1f;
        [SerializeField] private float widthMultiplier = 0.35f;
        [SerializeField] private AnimationCurve widthCurve = AnimationCurve.Linear(0f, 1f, 1f, 1f);
        [SerializeField] private Gradient colorGradient = new();
        [SerializeField] private bool generateLightingData = true;
        [SerializeField] private LineTextureMode textureMode = LineTextureMode.RepeatPerSegment;
        [SerializeField] private float uvSegmentLength = 0.5f;
        [SerializeField] private Material trailMaterial;
        [SerializeField] private MonoBehaviour wheelContactProviderBehaviour;
        [SerializeField] private Vector3 groundContactOffset = new Vector3(0f, 0.02f, 0f);

        private readonly List<TrailPoint> points = new();
        private static readonly List<Vector3> emptyNormals = new();
        private readonly object rebuildStateLock = new();
        private const int MinRenderablePointCount = 2;

        private Mesh trailMesh;
        private Vector3 lastEmitPosition;
        private bool hasLastEmitPosition;
        private bool shouldStartNewSegment = true;
        private bool meshHasRenderableData;
        private int rebuildGeneration;
        private int activeRebuildGeneration;
        private bool rebuildInProgress;
        private JobHandle pendingMeshJobHandle;
        private bool meshJobScheduled;
        private int scheduledMeshJobGeneration;

        private NativeArray<TrailPoint> nativeSnapshotPoints;
        private NativeArray<float> nativeSnapshotWidths;
        private NativeArray<Color> nativeSnapshotColors;
        private NativeArray<float> nativeSnapshotUCoordinates;
        private NativeArray<Vector3> nativeVertices;
        private NativeArray<Vector3> nativeNormals;
        private NativeArray<Color> nativeVertexColors;
        private NativeArray<Vector2> nativeUvs;
        private NativeArray<int> nativeIndices;
        private NativeArray<int> nativeMeshIndexCount;
        private NativeArray<byte> nativeMeshHasInvalidData;
        private NativeArray<Vector3> nativeNormalsEmpty;
        private int scheduledMeshPointCount;

        private IWheelContactProvider wheelContactProvider;

        public void Clear()
        {
            CompleteScheduledMeshJob();

            points.Clear();
            hasLastEmitPosition = false;
            shouldStartNewSegment = true;
            rebuildInProgress = false;
            rebuildGeneration += 1;
            activeRebuildGeneration = rebuildGeneration;
            trailMesh?.Clear();
            meshHasRenderableData = false;
        }

        private void Awake()
        {
            if (!ResolveWheelContactProvider())
            {
                Debug.LogError($"{nameof(WheelTrailRenderer)} on '{name}' could not find an {nameof(IWheelContactProvider)} on this object or children. Disabling component.", this);
                enabled = false;
                return;
            }

            trailMesh = new Mesh();
            const string MESH_NAME = "WheelTrailRendererMesh";
            trailMesh.name = MESH_NAME;
            trailMesh.MarkDynamic();
        }

        private void OnDisable()
        {
            Clear();
        }

        private void OnDestroy()
        {
            CompleteScheduledMeshJob();

            DisposeNativeMeshBuffers();

            if (trailMesh != null)
            {
                Destroy(trailMesh);
                trailMesh = null;
            }
        }

        private bool ResolveWheelContactProvider()
        {
            wheelContactProvider = null;

            if (wheelContactProviderBehaviour is IWheelContactProvider provider)
            {
                wheelContactProvider = provider;
                return true;
            }

            foreach (var monoBehaviour in GetComponentsInChildren<MonoBehaviour>(true))
            {
                if (monoBehaviour == null) continue;
                if (monoBehaviour == this) continue;
                if (monoBehaviour is not IWheelContactProvider contactProvider) continue;

                wheelContactProviderBehaviour = monoBehaviour;
                wheelContactProvider = contactProvider;
                return true;
            }

            return false;
        }

        private void LateUpdate()
        {
            FinishPendingMeshJob();
            CollectTrailData();
            RenderTrail();
        }

        private bool EmitPoint(Vector3 position, Vector3 normal)
        {
            if (!WheelTrailMath.IsFiniteVector3(position) || !WheelTrailMath.IsFiniteVector3(normal))
                return false;

            if (!hasLastEmitPosition)
            {
                AppendPoint(position, normal);
                return true;
            }

            float movedDistance = Vector3.Distance(lastEmitPosition, position);
            if (movedDistance < minVertexDistance)
                return false;

            AppendPoint(position, normal);
            return true;
        }

        private void AppendPoint(Vector3 position, Vector3 normal)
        {
            normal = WheelTrailMath.GetSafeNormalOrUp(normal);

            float distanceFromStart = 0f;
            int lastIndex = points.Count - 1;
            if (lastIndex >= 0)
            {
                distanceFromStart = points[lastIndex].distanceFromStart + Vector3.Distance(points[lastIndex].position, position);
                if (!float.IsFinite(distanceFromStart))
                    distanceFromStart = points[lastIndex].distanceFromStart;
            }

            points.Add(new TrailPoint
            {
                position = position,
                normal = normal,
                distanceFromStart = distanceFromStart,
                startsNewSegment = shouldStartNewSegment || points.Count == 0
            });

            lastEmitPosition = position;
            hasLastEmitPosition = true;
            shouldStartNewSegment = false;
        }

        private void RemovePointsBeyondMaxDistance()
        {
            if (points.Count == 0)
            {
                hasLastEmitPosition = false;
                shouldStartNewSegment = true;
                return;
            }

            float newestDistance = points[points.Count - 1].distanceFromStart;
            float oldestAllowedDistance = newestDistance - maxDistance;
            while (points.Count > 0 && points[0].distanceFromStart < oldestAllowedDistance)
                points.RemoveAt(0);

            hasLastEmitPosition = points.Count > 0;
            if (points.Count > 0)
            {
                // After trimming old points, force a seam so triangles are not bridged across removed history.
                TrailPoint firstPoint = points[0];
                firstPoint.startsNewSegment = true;
                points[0] = firstPoint;
            }
            else
                shouldStartNewSegment = true;
        }

        private void ScheduleMeshRebuild()
        {
            if (!CanStartRebuild())
                return;

            if (points.Count < MinRenderablePointCount)
            {
                trailMesh.Clear();
                meshHasRenderableData = false;
                return;
            }

            int pointCount = points.Count;
            int vertexCount = pointCount * 2;
            int maxIndexCount = (pointCount - 1) * 6;
            EnsureNativeSnapshotCapacity(pointCount);
            EnsureNativeJobMetadataCapacity();
            EnsureNativeWorkerCapacity(vertexCount, maxIndexCount, generateLightingData);
            PopulateSnapshotData(pointCount);

            int generation = BeginRebuildGeneration();

            WheelTrailMeshBuildJob wheelTrailMeshBuildJob = new WheelTrailMeshBuildJob
            {
                PointCount = pointCount,
                IncludeLightingData = generateLightingData,
                Points = nativeSnapshotPoints,
                Widths = nativeSnapshotWidths,
                Colors = nativeSnapshotColors,
                UCoordinates = nativeSnapshotUCoordinates,
                Vertices = nativeVertices,
                Normals = generateLightingData ? nativeNormals : nativeNormalsEmpty,
                OutColors = nativeVertexColors,
                Uvs = nativeUvs,
                Indices = nativeIndices,
                OutIndexCount = nativeMeshIndexCount,
                OutHasInvalidData = nativeMeshHasInvalidData
            };

            scheduledMeshJobGeneration = generation;
            scheduledMeshPointCount = pointCount;
            pendingMeshJobHandle = wheelTrailMeshBuildJob.Schedule();
            meshJobScheduled = true;
        }

        private void CompleteScheduledMeshJob()
        {
            if (!meshJobScheduled)
                return;

            pendingMeshJobHandle.Complete();
            meshJobScheduled = false;
        }

        private void ApplyCompletedMeshJob()
        {
            bool generationMismatch;
            lock (rebuildStateLock)
                generationMismatch = scheduledMeshJobGeneration != activeRebuildGeneration;

            try
            {
                // Drop late job results if a newer rebuild generation is already active.
                if (generationMismatch)
                    return;

                if (trailMesh == null)
                    return;

                bool hasInvalidData = nativeMeshHasInvalidData[0] != 0;
                int indexCount = nativeMeshIndexCount[0];
                bool hasRenderableData = indexCount > 0;
                int vertexCount = scheduledMeshPointCount * 2;

                if (hasInvalidData)
                {
                    trailMesh.Clear();
                    meshHasRenderableData = false;
                    return;
                }

                trailMesh.Clear();
                trailMesh.SetVertices(nativeVertices, 0, vertexCount);

                if (generateLightingData && nativeNormals.IsCreated && nativeNormals.Length >= vertexCount)
                    trailMesh.SetNormals(nativeNormals, 0, vertexCount);
                else
                    trailMesh.SetNormals(emptyNormals);

                trailMesh.SetColors(nativeVertexColors, 0, vertexCount);
                trailMesh.SetUVs(0, nativeUvs, 0, vertexCount);
                trailMesh.SetIndices(nativeIndices, 0, indexCount, MeshTopology.Triangles, 0, true);
                trailMesh.RecalculateBounds();
                meshHasRenderableData = hasRenderableData;
            }
            finally
            {
                lock (rebuildStateLock)
                    rebuildInProgress = false;
            }
        }

        private void DisposeNativeMeshBuffers()
        {
            DisposeIfCreated(ref nativeSnapshotPoints);
            DisposeIfCreated(ref nativeSnapshotWidths);
            DisposeIfCreated(ref nativeSnapshotColors);
            DisposeIfCreated(ref nativeSnapshotUCoordinates);
            DisposeIfCreated(ref nativeVertices);
            DisposeIfCreated(ref nativeNormals);
            DisposeIfCreated(ref nativeNormalsEmpty);
            DisposeIfCreated(ref nativeVertexColors);
            DisposeIfCreated(ref nativeUvs);
            DisposeIfCreated(ref nativeIndices);
            DisposeIfCreated(ref nativeMeshIndexCount);
            DisposeIfCreated(ref nativeMeshHasInvalidData);
        }

        private void EnsureNativeJobMetadataCapacity()
        {
            if (nativeMeshIndexCount.IsCreated)
                return;

            nativeMeshIndexCount = new NativeArray<int>(1, Allocator.Persistent);
            nativeMeshHasInvalidData = new NativeArray<byte>(1, Allocator.Persistent);
            nativeNormalsEmpty = new NativeArray<Vector3>(0, Allocator.Persistent);
        }

        private void EnsureNativeSnapshotCapacity(int pointCount)
        {
            int currentCapacity = nativeSnapshotPoints.IsCreated ? nativeSnapshotPoints.Length : 0;
            if (nativeSnapshotPoints.IsCreated && currentCapacity >= pointCount)
                return;

            int newCapacity = GetNextCapacity(currentCapacity, pointCount);
            DisposeIfCreated(ref nativeSnapshotPoints);
            DisposeIfCreated(ref nativeSnapshotWidths);
            DisposeIfCreated(ref nativeSnapshotColors);
            DisposeIfCreated(ref nativeSnapshotUCoordinates);

            nativeSnapshotPoints = new NativeArray<TrailPoint>(newCapacity, Allocator.Persistent);
            nativeSnapshotWidths = new NativeArray<float>(newCapacity, Allocator.Persistent);
            nativeSnapshotColors = new NativeArray<Color>(newCapacity, Allocator.Persistent);
            nativeSnapshotUCoordinates = new NativeArray<float>(newCapacity, Allocator.Persistent);
        }

        private void EnsureNativeWorkerCapacity(int vertexCount, int indexCount, bool includeLightingData)
        {
            int currentVertexCapacity = nativeVertices.IsCreated ? nativeVertices.Length : 0;
            if (!nativeVertices.IsCreated || currentVertexCapacity < vertexCount)
            {
                int newVertexCapacity = GetNextCapacity(currentVertexCapacity, vertexCount);
                DisposeIfCreated(ref nativeVertices);
                DisposeIfCreated(ref nativeVertexColors);
                DisposeIfCreated(ref nativeUvs);
                DisposeIfCreated(ref nativeNormals);

                nativeVertices = new NativeArray<Vector3>(newVertexCapacity, Allocator.Persistent);
                nativeVertexColors = new NativeArray<Color>(newVertexCapacity, Allocator.Persistent);
                nativeUvs = new NativeArray<Vector2>(newVertexCapacity, Allocator.Persistent);
                nativeNormals = new NativeArray<Vector3>(includeLightingData ? newVertexCapacity : 0, Allocator.Persistent);
            }
            else if (includeLightingData && (!nativeNormals.IsCreated || nativeNormals.Length < vertexCount))
            {
                DisposeIfCreated(ref nativeNormals);
                int newNormalCapacity = GetNextCapacity(0, vertexCount);
                nativeNormals = new NativeArray<Vector3>(newNormalCapacity, Allocator.Persistent);
            }
            else if (!includeLightingData && nativeNormals.IsCreated && nativeNormals.Length > 0)
            {
                DisposeIfCreated(ref nativeNormals);
                nativeNormals = new NativeArray<Vector3>(0, Allocator.Persistent);
            }

            int currentIndexCapacity = nativeIndices.IsCreated ? nativeIndices.Length : 0;
            if (!nativeIndices.IsCreated || currentIndexCapacity < indexCount)
            {
                int newIndexCapacity = GetNextCapacity(currentIndexCapacity, indexCount);
                DisposeIfCreated(ref nativeIndices);
                nativeIndices = new NativeArray<int>(newIndexCapacity, Allocator.Persistent);
            }
        }

        private static int GetNextCapacity(int currentCapacity, int requiredCapacity)
        {
            int nextCapacity = currentCapacity > 0 ? currentCapacity : 8;
            while (nextCapacity < requiredCapacity)
                nextCapacity *= 2;

            return nextCapacity;
        }

        private static void DisposeIfCreated<T>(ref NativeArray<T> array) where T : unmanaged
        {
            if (!array.IsCreated)
                return;

            array.Dispose();
        }

        private void FinishPendingMeshJob()
        {
            if (!meshJobScheduled)
                return;

            CompleteScheduledMeshJob();
            ApplyCompletedMeshJob();
        }

        private void CollectTrailData()
        {
            if (!wheelContactProvider.GetWheelContact(out Vector3 point, out Vector3 normal))
            {
                ResetTrailSegmentContinuity();
                return;
            }

            if (!EmitPoint(point + groundContactOffset, normal))
                return;

            RemovePointsBeyondMaxDistance();
            ScheduleMeshRebuild();
        }

        private void ResetTrailSegmentContinuity()
        {
            // Force a discontinuity when emission starts again.
            hasLastEmitPosition = false;
            shouldStartNewSegment = true;
        }

        private void RenderTrail()
        {
            if (!meshHasRenderableData)
                return;

            Graphics.DrawMesh(
                trailMesh,
                Matrix4x4.identity,
                trailMaterial,
                gameObject.layer,
                null,
                0,
                null,
                ShadowCastingMode.Off,
                false
            );
        }

        private bool CanStartRebuild()
        {
            lock (rebuildStateLock)
                return !rebuildInProgress && !meshJobScheduled;
        }

        private int BeginRebuildGeneration()
        {
            lock (rebuildStateLock)
            {
                rebuildGeneration += 1;
                activeRebuildGeneration = rebuildGeneration;
                rebuildInProgress = true;
                return activeRebuildGeneration;
            }
        }

        private void PopulateSnapshotData(int pointCount)
        {
            float newestDistance = points[pointCount - 1].distanceFromStart;
            float safeUvSegmentLength = Mathf.Max(WheelTrailMath.MinUvSegmentLength, uvSegmentLength);
            for (int pointIndex = 0; pointIndex < pointCount; pointIndex++)
            {
                TrailPoint point = points[pointIndex];
                nativeSnapshotPoints[pointIndex] = point;

                float age01 = ComputeAge01(newestDistance, point.distanceFromStart);
                float life01 = 1f - age01;
                float width = widthCurve.Evaluate(life01) * widthMultiplier;
                if (!float.IsFinite(width))
                    width = 0f;

                nativeSnapshotWidths[pointIndex] = width;
                nativeSnapshotColors[pointIndex] = colorGradient.Evaluate(age01);
                nativeSnapshotUCoordinates[pointIndex] = ComputeUCoordinate(pointIndex, point.distanceFromStart, newestDistance, safeUvSegmentLength);
            }
        }

        private float ComputeAge01(float newestDistance, float pointDistanceFromStart)
        {
            float distanceFromNewest = newestDistance - pointDistanceFromStart;
            if (maxDistance <= WheelTrailMath.MinDistanceEpsilon)
                return 1f;

            return Mathf.Clamp01(distanceFromNewest / maxDistance);
        }

        private float ComputeUCoordinate(int pointIndex, float pointDistanceFromStart, float newestDistance, float safeUvSegmentLength)
        {
            if (textureMode == LineTextureMode.RepeatPerSegment)
                return pointIndex;

            if (textureMode == LineTextureMode.Tile)
                return pointDistanceFromStart / safeUvSegmentLength;

            if (newestDistance <= WheelTrailMath.MinDistanceEpsilon)
                return 0f;

            return pointDistanceFromStart / newestDistance;
        }
    }
}

using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using UnityEngine;

namespace WheelTrail
{
    internal struct TrailPoint
    {
        public Vector3 position;
        public Vector3 normal;
        public float distanceFromStart;
        public bool startsNewSegment;
    }

    internal static class WheelTrailMath
    {
        public const float MinVectorSqrMagnitude = 0.000001f;
        public const float MinDistanceEpsilon = 0.0001f;
        public const float MinUvSegmentLength = 0.001f;

        public static bool IsFiniteVector3(Vector3 vector)
        {
            return float.IsFinite(vector.x) && float.IsFinite(vector.y) && float.IsFinite(vector.z);
        }

        public static Vector3 GetSafeNormalOrUp(Vector3 normal)
        {
            if (!IsFiniteVector3(normal) || normal.sqrMagnitude < MinVectorSqrMagnitude)
                return Vector3.up;

            return normal.normalized;
        }
    }

    [BurstCompile]
    internal struct WheelTrailMeshBuildJob : IJob
    {
        public int PointCount;
        public bool IncludeLightingData;

        [ReadOnly] public NativeArray<TrailPoint> Points;
        [ReadOnly] public NativeArray<float> Widths;
        [ReadOnly] public NativeArray<Color> Colors;
        [ReadOnly] public NativeArray<float> UCoordinates;

        [WriteOnly] public NativeArray<Vector3> Vertices;
        [WriteOnly] public NativeArray<Vector3> Normals;
        [WriteOnly] public NativeArray<Color> OutColors;
        [WriteOnly] public NativeArray<Vector2> Uvs;
        public NativeArray<int> Indices;

        public NativeArray<int> OutIndexCount;
        public NativeArray<byte> OutHasInvalidData;

        public void Execute()
        {
            OutHasInvalidData[0] = 0;

            int pointCount = PointCount;
            int vertexCount = pointCount * 2;

            if (pointCount < 2 || vertexCount <= 0)
            {
                OutIndexCount[0] = 0;
                return;
            }

            bool includeLightingData = IncludeLightingData;
            bool normalsWritable = includeLightingData && Normals.IsCreated && Normals.Length >= vertexCount;

            for (int pointIndex = 0; pointIndex < pointCount; pointIndex++)
            {
                TrailPoint point = Points[pointIndex];
                float halfWidth = Widths[pointIndex] * 0.5f;

                Vector3 normal = WheelTrailMath.GetSafeNormalOrUp(point.normal);
                Vector3 trailTangent = ComputeTrailTangent(Points, pointIndex, pointCount);
                Vector3 alongTrail = Vector3.ProjectOnPlane(trailTangent, normal);
                if (alongTrail.sqrMagnitude < WheelTrailMath.MinVectorSqrMagnitude)
                {
                    alongTrail = Vector3.ProjectOnPlane(Vector3.forward, normal);
                    if (alongTrail.sqrMagnitude < WheelTrailMath.MinVectorSqrMagnitude)
                        alongTrail = Vector3.ProjectOnPlane(Vector3.right, normal);
                }
                alongTrail.Normalize();

                Vector3 widthDirection = Vector3.Cross(normal, alongTrail);
                if (!WheelTrailMath.IsFiniteVector3(widthDirection) || widthDirection.sqrMagnitude < WheelTrailMath.MinVectorSqrMagnitude)
                    widthDirection = Vector3.right;
                else
                    widthDirection.Normalize();

                Vector3 leftWorld = point.position - widthDirection * halfWidth;
                Vector3 rightWorld = point.position + widthDirection * halfWidth;
                if (!WheelTrailMath.IsFiniteVector3(leftWorld) || !WheelTrailMath.IsFiniteVector3(rightWorld))
                {
                    OutHasInvalidData[0] = 1;
                    OutIndexCount[0] = 0;
                    return;
                }

                int vertexBaseIndex = pointIndex * 2;
                Vertices[vertexBaseIndex] = leftWorld;
                Vertices[vertexBaseIndex + 1] = rightWorld;
                OutColors[vertexBaseIndex] = Colors[pointIndex];
                OutColors[vertexBaseIndex + 1] = Colors[pointIndex];
                Uvs[vertexBaseIndex] = new Vector2(UCoordinates[pointIndex], 0f);
                Uvs[vertexBaseIndex + 1] = new Vector2(UCoordinates[pointIndex], 1f);

                if (normalsWritable)
                {
                    Normals[vertexBaseIndex] = normal;
                    Normals[vertexBaseIndex + 1] = normal;
                }
            }

            int indexWritePosition = 0;
            for (int pointIndex = 0; pointIndex < pointCount - 1; pointIndex++)
            {
                if (Points[pointIndex + 1].startsNewSegment)
                    continue;

                int baseVertex = pointIndex * 2;
                Indices[indexWritePosition] = baseVertex;
                Indices[indexWritePosition + 1] = baseVertex + 2;
                Indices[indexWritePosition + 2] = baseVertex + 1;
                Indices[indexWritePosition + 3] = baseVertex + 1;
                Indices[indexWritePosition + 4] = baseVertex + 2;
                Indices[indexWritePosition + 5] = baseVertex + 3;
                indexWritePosition += 6;
            }

            OutIndexCount[0] = indexWritePosition;
        }

        private static Vector3 ComputeTrailTangent(NativeArray<TrailPoint> points, int index, int pointCount)
        {
            if (pointCount < 2)
                return Vector3.forward;

            if (index == 0)
            {
                if (!points[1].startsNewSegment)
                    return points[1].position - points[0].position;
                return Vector3.forward;
            }

            if (index == pointCount - 1)
                return points[index].position - points[index - 1].position;

            if (!points[index + 1].startsNewSegment && !points[index].startsNewSegment)
                return points[index + 1].position - points[index - 1].position;

            if (!points[index + 1].startsNewSegment)
                return points[index + 1].position - points[index].position;

            if (!points[index].startsNewSegment)
                return points[index].position - points[index - 1].position;

            return points[index + 1].position - points[index].position;
        }
    }
}

using UnityEngine;

namespace WheelTrail
{
    [DisallowMultipleComponent]
    [RequireComponent(typeof(WheelCollider))]
    public class WheelColliderTrailContactProvider : MonoBehaviour, IWheelContactProvider
    {
        [SerializeField] private WheelCollider wheelCollider;

        private void Reset()
        {
            wheelCollider = GetComponent<WheelCollider>();
        }

        public bool GetWheelContact(out Vector3 point, out Vector3 normal)
        {
            if (wheelCollider == null)
            {
                point = default;
                normal = Vector3.up;
                return false;
            }

            bool isGrounded = wheelCollider.GetGroundHit(out WheelHit wheelHit);
            point = wheelHit.point;
            normal = wheelHit.normal;
            return isGrounded;
        }
    }
}

using UnityEngine;

namespace WheelTrail
{
    public interface IWheelContactProvider
    {
        bool GetWheelContact(out Vector3 point, out Vector3 normal);
    }
}