using UnityEngine;

namespace NBF
{
    /// <summary>
    /// Keeps ConfigurableJoint length close to a hard max by clamping only when overstretched.
    /// Joint still handles most of the simulation; this script only corrects overflow.
    /// </summary>
    [DefaultExecutionOrder(10000)]
    public class ConfigurableJointLengthClamp : MonoBehaviour
    {
        [Header("References")]
        [SerializeField] private ConfigurableJoint joint;
        [SerializeField] private Rigidbody body;

        [Header("Limit")]
        [Tooltip("<= 0 means use joint.linearLimit.limit")]
        [SerializeField] private float overrideMaxDistance = -1f;

        [Min(0f)]
        [SerializeField] private float tolerance = 0.002f;

        [Min(0f)]
        [SerializeField] private float maxCorrectionPerStep = 0.02f;

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

        [Header("Velocity")]
        [SerializeField] private bool removeSeparatingVelocity = true;

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

        [SerializeField] private bool estimateKinematicConnectedVelocity = true;

        private const float Epsilon = 1e-6f;

        private bool _hasPrevConnectedPoint;
        private Vector3 _prevConnectedPoint;

        private void Reset()
        {
            joint = GetComponent<ConfigurableJoint>();
            body = GetComponent<Rigidbody>();
        }

        private void OnEnable()
        {
            _hasPrevConnectedPoint = false;
        }

        private void OnValidate()
        {
            tolerance = Mathf.Max(0f, tolerance);
            maxCorrectionPerStep = Mathf.Max(0f, maxCorrectionPerStep);
            positionalCorrection = Mathf.Clamp01(positionalCorrection);
            radialVelocityDamping = Mathf.Clamp01(radialVelocityDamping);
        }

        private void FixedUpdate()
        {
            if (joint == null || body == null)
            {
                return;
            }

            float maxDistance = GetMaxDistance();
            if (maxDistance <= 0f)
            {
                return;
            }

            Vector3 bodyPoint = GetBodyAnchorWorldPoint();
            Vector3 connectedPoint = GetConnectedAnchorWorldPoint();

            Vector3 delta = bodyPoint - connectedPoint;
            float distance = delta.magnitude;
            if (distance < Epsilon)
            {
                CacheConnectedPoint(connectedPoint);
                return;
            }

            float exceed = distance - maxDistance;
            if (exceed <= tolerance)
            {
                CacheConnectedPoint(connectedPoint);
                return;
            }

            Vector3 dir = delta / distance;
            float correction = (exceed - tolerance) * positionalCorrection;
            if (maxCorrectionPerStep > 0f)
            {
                correction = Mathf.Min(correction, maxCorrectionPerStep);
            }

            if (correction > Epsilon && !body.isKinematic)
            {
                body.position -= dir * correction;
                body.WakeUp();
            }

            if (removeSeparatingVelocity && !body.isKinematic)
            {
                Vector3 connectedVelocity = GetConnectedVelocity(connectedPoint);
                Vector3 relative = body.linearVelocity - connectedVelocity;
                float separating = Vector3.Dot(relative, dir);
                if (separating > 0f)
                {
                    body.linearVelocity -= dir * (separating * radialVelocityDamping);
                }
            }

            CacheConnectedPoint(connectedPoint);
        }

        private float GetMaxDistance()
        {
            if (overrideMaxDistance > 0f)
            {
                return overrideMaxDistance;
            }

            return Mathf.Max(0f, joint.linearLimit.limit);
        }

        private Vector3 GetBodyAnchorWorldPoint()
        {
            return joint.transform.TransformPoint(joint.anchor);
        }

        private Vector3 GetConnectedAnchorWorldPoint()
        {
            Rigidbody connectedBody = joint.connectedBody;
            if (connectedBody != null)
            {
                return connectedBody.transform.TransformPoint(joint.connectedAnchor);
            }

            return joint.connectedAnchor;
        }

        private Vector3 GetConnectedVelocity(Vector3 connectedPoint)
        {
            Rigidbody connectedBody = joint.connectedBody;
            if (connectedBody == null)
            {
                return Vector3.zero;
            }

            if (!connectedBody.isKinematic)
            {
                return connectedBody.GetPointVelocity(connectedPoint);
            }

            if (!estimateKinematicConnectedVelocity)
            {
                return Vector3.zero;
            }

            if (!_hasPrevConnectedPoint)
            {
                return Vector3.zero;
            }

            float dt = Mathf.Max(Time.fixedDeltaTime, Epsilon);
            return (connectedPoint - _prevConnectedPoint) / dt;
        }

        private void CacheConnectedPoint(Vector3 connectedPoint)
        {
            _prevConnectedPoint = connectedPoint;
            _hasPrevConnectedPoint = true;
        }
    }
}
