using UnityEngine;

namespace NBF
{
    /// <summary>
    /// Max-distance rope constraint for two rigidbodies.
    /// It only corrects when distance is larger than maxDistance.
    /// </summary>
    [DefaultExecutionOrder(10000)]
    public class RigidbodyRopeConstraint : MonoBehaviour
    {
        private enum SolveMode
        {
            BodyBOnly = 0,
            MassWeighted = 1,
        }

        [Header("Bodies")]
        [SerializeField] private Rigidbody bodyA;
        [SerializeField] private Rigidbody bodyB;

        [Header("Optional Anchors")]
        [Tooltip("Optional. If null, uses bodyA.worldCenterOfMass.")]
        [SerializeField] private Transform anchorA;

        [Tooltip("Optional. If null, uses bodyB.worldCenterOfMass.")]
        [SerializeField] private Transform anchorB;

        [Header("Rope Limit")]
        [Min(0f)]
        [SerializeField] private float maxDistance = 0.5f;

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

        [Range(1, 12)]
        [SerializeField] private int solverIterations = 4;

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

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

        [SerializeField] private SolveMode solveMode = SolveMode.BodyBOnly;

        [Header("Velocity")]
        [Tooltip("When stretched, remove outward relative velocity.")]
        [SerializeField] private bool cancelSeparatingVelocity = true;

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

        [Tooltip("Estimate velocity for kinematic anchors from previous frame.")]
        [SerializeField] private bool estimateKinematicVelocity = true;

        private const float Epsilon = 1e-6f;

        private bool _hasPrevPointA;
        private bool _hasPrevPointB;
        private Vector3 _prevPointA;
        private Vector3 _prevPointB;

        private void OnEnable()
        {
            _hasPrevPointA = false;
            _hasPrevPointB = false;
        }

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

        private void FixedUpdate()
        {
            if (bodyA == null || bodyB == null)
            {
                return;
            }

            float dt = Mathf.Max(Time.fixedDeltaTime, Epsilon);

            Vector3 pointA = GetPointA();
            Vector3 pointB = GetPointB();

            Vector3 anchorVelA = GetAnchorVelocity(bodyA, pointA, ref _hasPrevPointA, ref _prevPointA, dt);
            Vector3 anchorVelB = GetAnchorVelocity(bodyB, pointB, ref _hasPrevPointB, ref _prevPointB, dt);

            for (int i = 0; i < solverIterations; i++)
            {
                if (!SolvePositionOnce())
                {
                    break;
                }
            }

            if (cancelSeparatingVelocity)
            {
                SolveVelocity(anchorVelA, anchorVelB);
            }
        }

        private bool SolvePositionOnce()
        {
            Vector3 pointA = GetPointA();
            Vector3 pointB = GetPointB();

            Vector3 delta = pointB - pointA;
            float distance = delta.magnitude;
            if (distance < Epsilon)
            {
                return false;
            }

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

            Vector3 dir = delta / distance;

            float correctionMagnitude = (exceed - tolerance) * positionalCorrection;
            if (maxCorrectionPerStep > 0f)
            {
                correctionMagnitude = Mathf.Min(correctionMagnitude, maxCorrectionPerStep);
            }

            if (correctionMagnitude <= Epsilon)
            {
                return false;
            }

            Vector3 correction = dir * correctionMagnitude;
            ApplyPositionCorrection(correction);
            return true;
        }

        private void SolveVelocity(Vector3 cachedAnchorVelA, Vector3 cachedAnchorVelB)
        {
            Vector3 pointA = GetPointA();
            Vector3 pointB = GetPointB();

            Vector3 delta = pointB - pointA;
            float distance = delta.magnitude;
            if (distance < Epsilon || distance <= maxDistance + tolerance)
            {
                return;
            }

            Vector3 dir = delta / distance;

            Vector3 velA = bodyA.isKinematic ? cachedAnchorVelA : bodyA.GetPointVelocity(pointA);
            Vector3 velB = bodyB.isKinematic ? cachedAnchorVelB : bodyB.GetPointVelocity(pointB);
            Vector3 relative = velB - velA;
            float separating = Vector3.Dot(relative, dir);

            if (separating <= 0f)
            {
                return;
            }

            float remove = separating * radialVelocityDamping;
            ApplyVelocityCorrection(dir, remove);
        }

        private void ApplyPositionCorrection(Vector3 correction)
        {
            bool aDynamic = !bodyA.isKinematic;
            bool bDynamic = !bodyB.isKinematic;

            if (!aDynamic && !bDynamic)
            {
                return;
            }

            if (solveMode == SolveMode.BodyBOnly)
            {
                if (bDynamic)
                {
                    bodyB.position -= correction;
                    bodyB.WakeUp();
                }
                else if (aDynamic)
                {
                    bodyA.position += correction;
                    bodyA.WakeUp();
                }

                return;
            }

            if (aDynamic && !bDynamic)
            {
                bodyA.position += correction;
                bodyA.WakeUp();
                return;
            }

            if (!aDynamic && bDynamic)
            {
                bodyB.position -= correction;
                bodyB.WakeUp();
                return;
            }

            float invMassA = GetInvMass(bodyA);
            float invMassB = GetInvMass(bodyB);
            float invMassSum = invMassA + invMassB;
            if (invMassSum < Epsilon)
            {
                return;
            }

            float ratioA = invMassA / invMassSum;
            float ratioB = invMassB / invMassSum;

            bodyA.position += correction * ratioA;
            bodyB.position -= correction * ratioB;

            bodyA.WakeUp();
            bodyB.WakeUp();
        }

        private void ApplyVelocityCorrection(Vector3 dir, float removingSpeed)
        {
            if (removingSpeed <= Epsilon)
            {
                return;
            }

            bool aDynamic = !bodyA.isKinematic;
            bool bDynamic = !bodyB.isKinematic;

            if (!aDynamic && !bDynamic)
            {
                return;
            }

            if (solveMode == SolveMode.BodyBOnly)
            {
                if (bDynamic)
                {
                    bodyB.linearVelocity -= dir * removingSpeed;
                }
                else if (aDynamic)
                {
                    bodyA.linearVelocity += dir * removingSpeed;
                }

                return;
            }

            if (aDynamic && !bDynamic)
            {
                bodyA.linearVelocity += dir * removingSpeed;
                return;
            }

            if (!aDynamic && bDynamic)
            {
                bodyB.linearVelocity -= dir * removingSpeed;
                return;
            }

            float invMassA = GetInvMass(bodyA);
            float invMassB = GetInvMass(bodyB);
            float invMassSum = invMassA + invMassB;
            if (invMassSum < Epsilon)
            {
                return;
            }

            float ratioA = invMassA / invMassSum;
            float ratioB = invMassB / invMassSum;

            bodyA.linearVelocity += dir * (removingSpeed * ratioA);
            bodyB.linearVelocity -= dir * (removingSpeed * ratioB);
        }

        private Vector3 GetPointA()
        {
            if (anchorA != null)
            {
                return anchorA.position;
            }

            return bodyA.worldCenterOfMass;
        }

        private Vector3 GetPointB()
        {
            if (anchorB != null)
            {
                return anchorB.position;
            }

            return bodyB.worldCenterOfMass;
        }

        private Vector3 GetAnchorVelocity(
            Rigidbody body,
            Vector3 point,
            ref bool hasPrev,
            ref Vector3 prevPoint,
            float dt)
        {
            if (body == null)
            {
                return Vector3.zero;
            }

            if (!body.isKinematic || !estimateKinematicVelocity)
            {
                hasPrev = true;
                prevPoint = point;
                return body.isKinematic ? Vector3.zero : body.GetPointVelocity(point);
            }

            Vector3 velocity = Vector3.zero;
            if (hasPrev)
            {
                velocity = (point - prevPoint) / dt;
            }

            hasPrev = true;
            prevPoint = point;
            return velocity;
        }

        private static float GetInvMass(Rigidbody body)
        {
            if (body == null || body.isKinematic)
            {
                return 0f;
            }

            float mass = Mathf.Max(body.mass, Epsilon);
            return 1f / mass;
        }
    }
}
