Spline Caps

using System;
using System.Linq;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Serialization;
using UnityEngine.Splines;
using Object = UnityEngine.Object;
#if UNITY_EDITOR
using UnityEditor;
#endif

namespace sc.modeling.splines.runtime.auxiliary
{
    [ExecuteAlways]
    public class SplineCaps : MonoBehaviour
    {
        public SplineContainer splineContainer;
        [SerializeField] [ HideInInspector]
        private int splineCount;

        [Serializable]
        public class Cap
        {
            public Cap(Position position)
            {
                this.position = position;
            }

            public enum Position
            {
                Start,
                End
            }
            public Position position;

            public GameObject prefab;
            [SerializeField] [ HideInInspector]
            private int previousPrefabID;

            public bool HasPrefabChanged()
            {
                if (prefab == null) return true;

                int hashCode = prefab.GetHashCode();
                if (hashCode != previousPrefabID)
                {
                    previousPrefabID = hashCode;

                    return true;
                }

                return false;
            }

            [Space]
            [Header("Position")]
            [Tooltip("Positional offset, relative to the curve's tangent")]
            public Vector3 offset;
            [Tooltip("Shifts the object along the spline curve by this many units")]
            [Min(0f)]
            public float shift = 0f;

            [Space]
            [Header("Rotation")]
            [Tooltip("Align the object's forward direction to the tangent of the spline")]
            [FormerlySerializedAs("alignRotation")]
            public bool align = true;
            [Tooltip("Rotation in degrees, added to the object's rotation")]
            [FormerlySerializedAs("addedRotation")]
            public Vector3 rotation;

            [Header("Scale")]
            public Vector3 scale = Vector3.one;

            //Save a reference to the instantiated objects, so they can be accessed again, deleted when necessary.
            [HideInInspector]
            public GameObject[] instances = Array.Empty<GameObject>();

            public bool HasMissingInstances()
            {
                for (int i = 0; i < instances.Length; i++)
                {
                    if (instances[i] == null) return true;
                }
                return false;
            }
        }

        [Space]

        public Cap[] caps = new[]
        {
            new Cap(Cap.Position.Start),
            new Cap(Cap.Position.End)
        };

        [Space]

        [Tooltip("Hide the prefab instances in the hierarchy")]
        public bool hideInstances = false;

        private void Reset()
        {
            splineContainer = gameObject.GetComponentInParent<SplineContainer>();
        }

        private void OnEnable()
        {
            Spline.Changed += OnSplineChanged;
        }

        private void OnDisable()
        {
            Spline.Changed -= OnSplineChanged;
        }

        private void OnSplineChanged(Spline spline, int index, SplineModification modificationType)
        {
            if (!splineContainer) return;

            //Spline belongs to the assigned container?
            var splineIndex = Array.IndexOf(splineContainer.Splines.ToArray(), spline);
            if (splineIndex < 0)
                return;

            Apply();
        }

        public void Apply()
        {
            if (!splineContainer) return;

            var splineCountChanged = splineContainer.Splines.Count != splineCount;
            //if(splineCountChanged) Debug.Log($"Spline count changed from {splineCount} to {splineContainer.Splines.Count}");
            splineCount = splineContainer.Splines.Count;

            for (int i = 0; i < caps.Length; i++)
            {
                if (splineCountChanged || caps[i].HasPrefabChanged() || caps[i].HasMissingInstances())
                {
                    //Debug.Log($"Respawning cap {i}");
                    Respawn(caps[i]);
                }
            }

            for (int i = 0; i < caps.Length; i++)
            {
                ApplyTransform(caps[i]);
            }
        }

        private void Respawn(Cap cap)
        {
            //Destroy any existing instances
            for (int i = 0; i < cap.instances.Length; i++)
            {
                CoreUtils.Destroy(cap.instances[i]);
            }
            cap.instances = Array.Empty<GameObject>();

            if (cap.prefab == null) return;

            //Respawn the prefabs (once per spline)
            cap.instances = new GameObject[splineContainer.Splines.Count];

            for (int i = 0; i < cap.instances.Length; i++)
            {
                GameObject instance = Instantiate(cap.prefab);
                instance.transform.SetParent(this.transform);

                cap.instances[i] = instance;
            }
        }

        void ApplyTransform(Cap cap)
        {
            for (int splineIndex = 0; splineIndex < cap.instances.Length; splineIndex++)
            {
                TransformToSpline(cap, cap.instances[splineIndex].transform, splineIndex);

                //Gray out fields as any chances would be overwritten anyway
                cap.instances[splineIndex].transform.hideFlags =  HideFlags.NotEditable;

                cap.instances[splineIndex].hideFlags = hideInstances ? HideFlags.HideInHierarchy : HideFlags.None;
            }
        }

        private bool IsPrefab(Object source)
        {
            bool isPrefab = false;

            #if UNITY_EDITOR
            if (UnityEditor.PrefabUtility.GetPrefabAssetType(source) == PrefabAssetType.Variant)
            {
                //PrefabUtility.GetCorrespondingObjectFromSource still returns the base prefab. However, this does work.
                isPrefab = source;
            }
            else
            {
                isPrefab = UnityEditor.PrefabUtility.GetCorrespondingObjectFromOriginalSource(source);
            }
            #endif

            return isPrefab;
        }

        private new GameObject Instantiate(Object source)
        {
            bool isPrefab = IsPrefab(source);
            GameObject instance = null;

            #if UNITY_EDITOR
            if (isPrefab)
            {
                instance = UnityEditor.PrefabUtility.InstantiatePrefab(source, this.gameObject.scene) as GameObject;
            }
            #endif

            //Non-prefabs and builds
            if (!isPrefab)
            {
                instance = GameObject.Instantiate(source) as GameObject;
            }

            return instance;
        }

        private void TransformToSpline(Cap cap, Transform target, int splineIndex)
        {
            //Coincidentally the two enum values correspond to 0 and 1
            float t = (int)cap.position;

            //Shift along spline by X-units
            float shift = cap.shift / splineContainer.Splines[splineIndex].GetLength();
            if (cap.position == Cap.Position.End) shift = -shift;

            t += shift;

            //Ensure a tangent can always be derived at very the start and end
            t = Mathf.Clamp(t, 0.0001f, 0.9999f);

            splineContainer.Splines[splineIndex].Evaluate(t, out float3 splinePoint, out float3 tangent, out float3 up);

            float3 position = splinePoint;
            float3 forward = math.normalize(tangent);
            float3 right = math.cross(forward, up);

            //Offset
            position += right * cap.offset.x;
            position += up * cap.offset.y;
            position += forward * cap.offset.z;

            //Position of point on spline in world-space
            position = splineContainer.transform.TransformPoint(position);

            Quaternion rotation = Quaternion.Euler(cap.rotation);
            if (cap.align)
            {
                rotation = Quaternion.LookRotation(forward) * Quaternion.Euler(cap.rotation);
            }

            target.SetPositionAndRotation(position, rotation);
            target.localScale = cap.scale;
        }
    }

    #if UNITY_EDITOR
    [CustomEditor(typeof(SplineCaps))]
    [CanEditMultipleObjects]
    public class SplineCapsEditor : Editor
    {
        public override void OnInspectorGUI()
        {
            EditorGUI.BeginChangeCheck();

            base.OnInspectorGUI();

            if (EditorGUI.EndChangeCheck())
            {
                foreach (var m_target in targets)
                    ((SplineCaps)m_target).Apply();
            }
        }
    }
    #endif
}