using System;
using System.Collections.Generic;
using LibMMD.Model;
using LibMMD.Util;
using UnityEngine;
using Tools = LibMMD.Util.Tools;

namespace LibMMD.Motion
{
    internal sealed class MmdIkSolver
    {
        private readonly MmdModel Model;
        private readonly BoneImage[] BoneImages;
        private readonly Action<BoneImage> UpdateLocalMatrix;
        private readonly Action<int> UpdateBoneSelfTransform;
        private readonly Dictionary<int, Vector3> _ikPoleVectors = new Dictionary<int, Vector3>();

        internal MmdIkSolver(MmdModel model, BoneImage[] boneImages,
            Action<BoneImage> updateLocalMatrix, Action<int> updateBoneSelfTransform)
        {
            Model = model ?? throw new ArgumentNullException(nameof(model));
            BoneImages = boneImages ?? throw new ArgumentNullException(nameof(boneImages));
            UpdateLocalMatrix = updateLocalMatrix ?? throw new ArgumentNullException(nameof(updateLocalMatrix));
            UpdateBoneSelfTransform = updateBoneSelfTransform ?? throw new ArgumentNullException(nameof(updateBoneSelfTransform));
        }

        internal void Solve(int index)
        {
            var image = BoneImages[index];
            if (!image.HasIk || !image.IkEnabled) return;
            var ikLinkNum = image.IkLinks.Length;
            for (var i = 0; i < ikLinkNum; ++i)
            {
                BoneImages[image.IkLinks[i]].IkRotation = Quaternion.identity;
            }
            var ikPosition = MathUtil.GetTransFromMatrix4X4(image.LocalMatrix);
            var targetPosition = MathUtil.GetTransFromMatrix4X4(BoneImages[image.IkTarget].LocalMatrix);
            var targetRotationBeforeIk = BoneImages[image.IkTarget].LocalMatrix.rotation;
            var ikError = ikPosition - targetPosition;
            var isTwoBoneLeg = _ikPoleVectors.ContainsKey(index);

            // Non-leg chains can return immediately. A recognized leg must pass through Soft
            // IK even when its unmodified target is already reached: exact full extension is
            // the singular pose that Soft IK is intended to avoid.
            if (!isTwoBoneLeg && Vector3.Dot(ikError, ikError) < Tools.MmdMathConstEps)
            {
                SetIkTargetRotation(image, targetRotationBeforeIk);
                UpdateLocalMatrix(BoneImages[image.IkTarget]);
                return;
            }
            // Keep every recognized two-bone leg away from the fully extended/folded
            // singularities. PMD has less link metadata than PMX, but ConfigureDefaultFootIkPoles
            // has already positively identified the same thigh/knee/ankle chain.
            if (isTwoBoneLeg)
                ikPosition = ApplySoftLegReach(image, ikPosition);
            ikError = ikPosition - targetPosition;

            if (Vector3.Dot(ikError, ikError) < Tools.MmdMathConstEps)
            {
                SetIkTargetRotation(image, targetRotationBeforeIk);
                UpdateLocalMatrix(BoneImages[image.IkTarget]);
                return;
            }

            // A two-bone leg has a closed-form solution once its bend pole is known. CCD is
            // under-determined near full extension and, on real VMD interpolation boundaries,
            // can select a different intermediate solution for one frame before returning on
            // the next. Use the unique pole-oriented solution for recognized legs.
            if (isTwoBoneLeg && TrySolveTwoBoneLeg(image, ikPosition, _ikPoleVectors[index]))
            {
                FinalizePoleAndTargetRotation(index, image, targetRotationBeforeIk);
                return;
            }

            var ikt = image.CcdIterateLimit / 2;
            for (var i = 0; i < image.CcdIterateLimit; ++i)
            {
                for (var j = 0; j < ikLinkNum; ++j)
                {
                    if (image.IkFixTypes[j] == BoneImage.AxisFixType.FixAll) continue;
                    var ikImage = BoneImages[image.IkLinks[j]];
                    var ikLinkPosition = MathUtil.GetTransFromMatrix4X4(ikImage.LocalMatrix);
                    var targetDirection = ikLinkPosition - targetPosition;
                    var ikDirection = ikLinkPosition - ikPosition;

                    targetDirection.Normalize();
                    ikDirection.Normalize();

                    var ikRotateAxis = Vector3.Cross(targetDirection, ikDirection);
                    if (ikRotateAxis.sqrMagnitude < 1e-8f)
                        continue;

                    var localizationMatrix = ikImage.HasParent ? BoneImages[ikImage.Parent].LocalMatrix : Matrix4x4.identity;
                    var skipLink = false;
                    // A one-axis IK link (normally a knee) must remain a hinge during every
                    // CCD iteration. Releasing the axis during the latter half and clamping
                    // only afterwards lets Y/Z rotation enter the solution, causing the knee
                    // plane to wander or rebound between adjacent frames.
                    if (image.IkLinkLimited[j] && image.IkFixTypes[j] != BoneImage.AxisFixType.FixNone)
                    {
                        switch (image.IkFixTypes[j])
                        {
                            case BoneImage.AxisFixType.FixX:
                            {
                                var dotVal = Vector3.Dot(ikRotateAxis, MathUtil.Matrix4x4ColDowngrade(localizationMatrix, 0));
                                if (Mathf.Abs(dotVal) < Tools.MmdMathConstEps) { skipLink = true; break; }
                                ikRotateAxis.x = Nabs(dotVal);
                                ikRotateAxis.y = ikRotateAxis.z = 0.0f;
                                break;
                            }
                            case BoneImage.AxisFixType.FixY:
                            {
                                var dotVal = Vector3.Dot(ikRotateAxis, MathUtil.Matrix4x4ColDowngrade(localizationMatrix, 1));
                                if (Mathf.Abs(dotVal) < Tools.MmdMathConstEps) { skipLink = true; break; }
                                ikRotateAxis.y = Nabs(dotVal);
                                ikRotateAxis.x = ikRotateAxis.z = 0.0f;
                                break;
                            }
                            case BoneImage.AxisFixType.FixZ:
                            {
                                var dotVal = Vector3.Dot(ikRotateAxis, MathUtil.Matrix4x4ColDowngrade(localizationMatrix, 2));
                                if (Mathf.Abs(dotVal) < Tools.MmdMathConstEps) { skipLink = true; break; }
                                ikRotateAxis.z = Nabs(dotVal);
                                ikRotateAxis.x = ikRotateAxis.y = 0.0f;
                                break;
                            }
                        }
                    }
                    else
                    {
                        ikRotateAxis = Matrix4x4.Transpose(localizationMatrix).MultiplyVector(ikRotateAxis);
                        ikRotateAxis.Normalize();
                        if (ikRotateAxis.sqrMagnitude < Tools.MmdMathConstEps)
                            skipLink = true;
                    }
                    if (skipLink) continue;
                    var ikRotateAngle =
                        Mathf.Min(Mathf.Acos(Mathf.Clamp(Vector3.Dot(targetDirection, ikDirection), -1.0f, 1.0f)),
                            image.CcdAngleLimit * (j + 1));
                    ikImage.IkRotation =
                        Quaternion.AngleAxis((float) (ikRotateAngle / Math.PI * 180.0), ikRotateAxis) * ikImage.IkRotation;
                    if (image.IkLinkLimited[j])
                    {
                        var localRotation = ikImage.IkRotation * ikImage.PreIkRotation;
                        switch (image.IkTransformOrders[j])
                        {
                            case BoneImage.AxisTransformOrder.OrderZxy:
                            {
                                var eularAngle = MathUtil.QuaternionToZxy(localRotation);
                                eularAngle = LimitEularAngle(eularAngle, image.IkLinkLimitsMin[j],
                                    image.IkLinkLimitsMax[j], image.IkFixTypes[j] != BoneImage.AxisFixType.FixX && i < ikt);
                                localRotation = MathUtil.ZxyToQuaternion(eularAngle);
                                break;
                            }
                            case BoneImage.AxisTransformOrder.OrderXyz:
                            {
                                var eularAngle = MathUtil.QuaternionToXyz(localRotation);
                                eularAngle = LimitEularAngle(eularAngle, image.IkLinkLimitsMin[j],
                                    image.IkLinkLimitsMax[j], image.IkFixTypes[j] != BoneImage.AxisFixType.FixX && i < ikt);
                                localRotation = MathUtil.XyzToQuaternion(eularAngle);
                                break;
                            }
                            case BoneImage.AxisTransformOrder.OrderYzx:
                            {
                                var eularAngle = MathUtil.QuaternionToYzx(localRotation);
                                eularAngle = LimitEularAngle(eularAngle, image.IkLinkLimitsMin[j],
                                    image.IkLinkLimitsMax[j], image.IkFixTypes[j] != BoneImage.AxisFixType.FixX && i < ikt);
                                localRotation = MathUtil.YzxToQuaternion(eularAngle);
                                break;
                            }
                            default:
                                throw new ArgumentOutOfRangeException();
                        }
                        // localRotation belongs to this IK link. Using the controller bone's
                        // PreIkRotation here leaks foot-IK rotation into the knee, which only
                        // became visible once IK was enabled during baking.
                        ikImage.IkRotation = localRotation * Quaternion.Inverse(ikImage.PreIkRotation);
                    }
                    for (var k = 0; k <= j; ++k)
                    {
                        var linkImage = BoneImages[image.IkLinks[j - k]];
                        linkImage.TotalRotation = linkImage.IkRotation * linkImage.PreIkRotation;
                        UpdateLocalMatrix(linkImage);
                    }
                    UpdateBoneSelfTransform(image.IkTarget);
                    targetPosition = MathUtil.Matrix4x4ColDowngrade(BoneImages[image.IkTarget].LocalMatrix, 3);
                }
                ikError = ikPosition - targetPosition;
                if (Vector3.Dot(ikError, ikError) < Tools.MmdMathConstEps)
                {
                    FinalizePoleAndTargetRotation(index, image, targetRotationBeforeIk);
                    return;
                }
            }
            FinalizePoleAndTargetRotation(index, image, targetRotationBeforeIk);
        }

        private void FinalizePoleAndTargetRotation(int ikBoneIndex, BoneImage image,
            Quaternion targetRotationBeforeIk)
        {
            if (!_ikPoleVectors.TryGetValue(ikBoneIndex, out var poleDirection)) return;

            ApplyPoleCorrection(image, poleDirection);
            SetIkTargetRotation(image, targetRotationBeforeIk);
            UpdateLocalMatrix(BoneImages[image.IkTarget]);
        }

        private Vector3 ApplySoftLegReach(BoneImage ikImage, Vector3 desiredPosition)
        {
            var thighIndex = -1;
            var kneeIndex = -1;
            for (var j = 0; j < ikImage.IkLinks.Length; ++j)
            {
                var linkIndex = ikImage.IkLinks[j];
                if (ikImage.IkFixTypes[j] == BoneImage.AxisFixType.FixNone) thighIndex = linkIndex;
                if (kneeIndex < 0 && ikImage.IkFixTypes[j] == BoneImage.AxisFixType.FixX) kneeIndex = linkIndex;
            }
            if (thighIndex < 0 || kneeIndex < 0) return desiredPosition;

            var ankleIndex = ikImage.IkTarget;
            var upperLength = (Model.Bones[kneeIndex].Position - Model.Bones[thighIndex].Position).magnitude;
            var lowerLength = (Model.Bones[ankleIndex].Position - Model.Bones[kneeIndex].Position).magnitude;
            if (upperLength < 1e-6f || lowerLength < 1e-6f) return desiredPosition;

            var hipPosition = MathUtil.GetTransFromMatrix4X4(BoneImages[thighIndex].LocalMatrix);
            var direction = desiredPosition - hipPosition;
            var distance = direction.magnitude;
            var totalLength = upperLength + lowerLength;
            // Leave enough bend for the pole plane to remain measurable even when the IK
            // control is well outside the reachable sphere. 0.1% still made the knee's
            // projected pole small enough for frame-to-frame float noise to become visible
            // as trembling on long-legged models.
            var margin = Mathf.Max(totalLength * 0.005f, 1e-4f);
            var minimum = Mathf.Abs(upperLength - lowerLength) + margin;
            var maximum = Mathf.Max(minimum, totalLength - margin);

            // Inside the minimum reach sphere the desired direction becomes numerically
            // unstable and reverses as the source-model target crosses the hip. Preserve the
            // current leg direction there; only correct the unreachable distance.
            if (distance < minimum)
            {
                direction = MathUtil.GetTransFromMatrix4X4(BoneImages[ankleIndex].LocalMatrix) - hipPosition;
                if (direction.sqrMagnitude < 1e-8f) direction = Vector3.down;
                direction.Normalize();
            }
            else
            {
                direction /= distance;
            }

            // Ease into the reach limit instead of hard-clamping at it. A hard clamp has a
            // derivative discontinuity and shows up as a one-frame knee/hip pop. This curve
            // is position- and slope-continuous at softStart, then approaches maximum without
            // ever reaching the straight-leg singularity.
            var softRange = Mathf.Max(totalLength * 0.01f, margin * 2f);
            softRange = Mathf.Min(softRange, Mathf.Max((maximum - minimum) * 0.5f, margin));
            var softStart = maximum - softRange;
            float solvedDistance;
            if (distance > softStart)
            {
                var softOffset = distance - softStart;
                solvedDistance = softStart + softRange *
                    (1f - Mathf.Exp(-softOffset / softRange));
            }
            else
            {
                solvedDistance = distance;
            }
            solvedDistance = Mathf.Clamp(solvedDistance, minimum, maximum);
            return hipPosition + direction * solvedDistance;
        }

        private bool TrySolveTwoBoneLeg(BoneImage ikImage, Vector3 solvedAnklePosition,
            Vector3 poleDirection)
        {
            var thighIndex = -1;
            var kneeIndex = -1;
            for (var j = 0; j < ikImage.IkLinks.Length; ++j)
            {
                if (ikImage.IkFixTypes[j] == BoneImage.AxisFixType.FixNone)
                    thighIndex = ikImage.IkLinks[j];
                else if (kneeIndex < 0 && ikImage.IkFixTypes[j] == BoneImage.AxisFixType.FixX)
                    kneeIndex = ikImage.IkLinks[j];
            }
            if (thighIndex < 0 || kneeIndex < 0) return false;

            var thigh = BoneImages[thighIndex];
            var knee = BoneImages[kneeIndex];
            var ankle = BoneImages[ikImage.IkTarget];
            var hipPosition = MathUtil.GetTransFromMatrix4X4(thigh.LocalMatrix);
            var kneePosition = MathUtil.GetTransFromMatrix4X4(knee.LocalMatrix);
            var anklePosition = MathUtil.GetTransFromMatrix4X4(ankle.LocalMatrix);
            var upper = kneePosition - hipPosition;
            var lower = anklePosition - kneePosition;
            var toAnkle = solvedAnklePosition - hipPosition;
            var upperLength = upper.magnitude;
            var lowerLength = lower.magnitude;
            var distance = toAnkle.magnitude;
            if (upperLength < 1e-6f || lowerLength < 1e-6f || distance < 1e-6f)
                return false;

            var axis = toAnkle / distance;
            var parentRotation = thigh.HasParent
                ? BoneImages[thigh.Parent].LocalMatrix.rotation
                : Quaternion.identity;
            var pole = parentRotation * poleDirection;
            pole -= Vector3.Dot(pole, axis) * axis;
            if (pole.sqrMagnitude < 1e-8f)
            {
                pole = upper - Vector3.Dot(upper, axis) * axis;
                if (pole.sqrMagnitude < 1e-8f) return false;
            }
            pole.Normalize();

            var along = (upperLength * upperLength - lowerLength * lowerLength +
                         distance * distance) / (2f * distance);
            var height = Mathf.Sqrt(Mathf.Max(0f,
                upperLength * upperLength - along * along));
            var solvedKneePosition = hipPosition + axis * along + pole * height;
            var solvedUpper = solvedKneePosition - hipPosition;
            var solvedLower = solvedAnklePosition - solvedKneePosition;

            var thighWorldBefore = thigh.LocalMatrix.rotation;
            var kneeWorldBefore = knee.LocalMatrix.rotation;
            var thighDelta = Quaternion.FromToRotation(upper, solvedUpper);
            var thighWorld = thighDelta * thighWorldBefore;
            var thighLocal = Quaternion.Inverse(parentRotation) * thighWorld;
            thigh.IkRotation = thighLocal * Quaternion.Inverse(thigh.PreIkRotation);
            thigh.TotalRotation = thigh.IkRotation * thigh.PreIkRotation;
            UpdateLocalMatrix(thigh);

            var rotatedLower = thighDelta * lower;
            var kneeDelta = Quaternion.FromToRotation(rotatedLower, solvedLower);
            var kneeWorld = kneeDelta * thighDelta * kneeWorldBefore;
            var kneeLocal = Quaternion.Inverse(thighWorld) * kneeWorld;
            knee.IkRotation = kneeLocal * Quaternion.Inverse(knee.PreIkRotation);
            knee.TotalRotation = knee.IkRotation * knee.PreIkRotation;
            UpdateLocalMatrix(knee);
            UpdateBoneSelfTransform(ikImage.IkTarget);
            return true;
        }

        // CCD reaches the target position but has no preferred bend plane. Align the knee's
        // projection around the hip-to-ankle axis with the configured pole without changing
        // the solved leg length. Near a straight leg the projection is undefined, so retaining
        // the CCD result is safer than allowing a 180-degree knee flip.
        private void ApplyPoleCorrection(BoneImage ikImage, Vector3 poleDirection)
        {
            var thighLink = -1;
            var kneeLink = -1;
            for (var j = 0; j < ikImage.IkLinks.Length; ++j)
            {
                if (ikImage.IkFixTypes[j] == BoneImage.AxisFixType.FixNone) thighLink = j;
                if (kneeLink < 0 && ikImage.IkFixTypes[j] == BoneImage.AxisFixType.FixX) kneeLink = j;
            }
            if (thighLink < 0 || kneeLink < 0) return;

            var thighIndex = ikImage.IkLinks[thighLink];
            var kneeIndex = ikImage.IkLinks[kneeLink];
            var thigh = BoneImages[thighIndex];
            var knee = BoneImages[kneeIndex];
            var hipPosition = MathUtil.GetTransFromMatrix4X4(thigh.LocalMatrix);
            var kneePosition = MathUtil.GetTransFromMatrix4X4(knee.LocalMatrix);
            var anklePosition = MathUtil.GetTransFromMatrix4X4(BoneImages[ikImage.IkTarget].LocalMatrix);
            var axis = anklePosition - hipPosition;
            if (axis.sqrMagnitude < 1e-8f) return;
            axis.Normalize();

            // Pole vectors are defined relative to the leg parent. Keeping the pole in world
            // space makes a turned pelvis pull both knees toward the old global direction,
            // which appears as an unnecessary inward thigh/kneecap twist.
            var parentRotation = thigh.HasParent
                ? BoneImages[thigh.Parent].LocalMatrix.rotation
                : Quaternion.identity;
            // IK replaces the authored FK thigh swing. Letting that discarded rotation rotate
            // the pole makes the preferred plane move every frame, producing both pigeon-toed
            // and bow-legged knees. This applies equally to PMD and PMX; their ankle-rotation
            // handling remains separate after the positional solve.
            poleDirection = parentRotation * poleDirection;

            var currentPole = kneePosition - hipPosition;
            currentPole -= Vector3.Dot(currentPole, axis) * axis;
            var desiredPole = poleDirection - Vector3.Dot(poleDirection, axis) * axis;
            var legLengthSquared = (Model.Bones[kneeIndex].Position - Model.Bones[thighIndex].Position).sqrMagnitude +
                                   (Model.Bones[ikImage.IkTarget].Position - Model.Bones[kneeIndex].Position).sqrMagnitude;
            var poleEpsilon = Mathf.Max(legLengthSquared * 1e-8f, 1e-10f);

            if (currentPole.sqrMagnitude < poleEpsilon || desiredPole.sqrMagnitude < poleEpsilon) return;
            currentPole.Normalize();
            desiredPole.Normalize();

            // Rotate the solved chain all the way onto its pole plane. This rotation is about
            // the hip-to-ankle axis, so it does not change either segment length or the solved
            // ankle position. The former per-frame angle cap left a model-dependent part of
            // CCD's arbitrary bend plane in the result: deep poses could retain an inward
            // knee, while almost-straight poses amplified tiny target changes into trembling.
            // Clamp only the numerical endpoint; SignedAngle already returns [-180, 180].
            var correctionAngle = Mathf.Clamp(
                Vector3.SignedAngle(currentPole, desiredPole, axis), -180f, 180f);
            var correction = Quaternion.AngleAxis(correctionAngle, axis);
            var localCorrection = Quaternion.Inverse(parentRotation) * correction * parentRotation;
            thigh.IkRotation = localCorrection * thigh.IkRotation;
            thigh.TotalRotation = thigh.IkRotation * thigh.PreIkRotation;
            UpdateLocalMatrix(thigh);
            knee.TotalRotation = knee.IkRotation * knee.PreIkRotation;
            UpdateLocalMatrix(knee);
            UpdateBoneSelfTransform(ikImage.IkTarget);
        }

        private static float Nabs(float x)
        {
            if (x > 0.0f)
            {
                return 1.0f;
            }
            return -1.0f;
        }

        private static Vector3 LimitEularAngle(Vector3 eular, Vector3 eularMin, Vector3 eularMax, bool ikt)
        {
            var result = eular;
            for (var i = 0; i < 3; ++i)
            {
                if (result[i] < eularMin[i])
                {
                    var tf = 2 * eularMin[i] - result[i];
                    if (tf <= eularMax[i] && ikt)
                    {
                        result[i] = tf;
                    }
                    else
                    {
                        result[i] = eularMin[i];
                    }
                }
                if (result[i] > eularMax[i])
                {
                    var tf = 2 * eularMax[i] - result[i];
                    if (tf >= eularMin[i] && ikt)
                    {
                        result[i] = tf;
                    }
                    else
                    {
                        result[i] = eularMax[i];
                    }
                }
            }
            return result;
        }

        internal void SetPoleVector(int ikBoneIndex, Vector3 poleDirection)
        {
            if (poleDirection == Vector3.zero)
                _ikPoleVectors.Remove(ikBoneIndex);
            else
                _ikPoleVectors[ikBoneIndex] = poleDirection;
        }

        // Preserve the ankle's authored pre-IK world orientation so lower-body facing (for
        // example Rabbit Hole's initial 180-degree turn) is retained. If the VMD also authors
        // a foot-IK rotation, apply that controller-local rotation as an additional world
        // delta rather than replacing the facing-aware ankle orientation.
        private void SetIkTargetRotation(BoneImage ikBone, Quaternion targetRotationBeforeIk)
        {
            var tgt = BoneImages[ikBone.IkTarget];
            var ikParentRotation = ikBone.HasParent
                ? BoneImages[ikBone.Parent].LocalMatrix.rotation
                : Quaternion.identity;
            var targetParentRotation = tgt.HasParent
                ? BoneImages[tgt.Parent].LocalMatrix.rotation
                : Quaternion.identity;

            var worldDelta = ikParentRotation * ikBone.Rotation * Quaternion.Inverse(ikParentRotation);
            var targetWorldRotation = worldDelta * targetRotationBeforeIk;

            // IK output is transient. Writing it into Rotation makes an unanimated ankle use
            // the previous frame's corrected pose as the next frame's authored input, so even
            // a tiny correction accumulates into a slow continuous roll. TotalRotation drives
            // this frame's matrix while Rotation remains owned by VMD/FK.
            tgt.TotalRotation = Quaternion.Inverse(targetParentRotation) * targetWorldRotation;
        }
    }
}