using UnityEngine;

public class SolverCCD
{
    public RobotJoint[] Joints;
    public Transform EndEffector;

    public int MaxIterations = 40;
    public float PositionTolerance = 0.001f;
    public float OrientationToleranceDeg = 3f;

    public float MaxStepDegrees = 2f;
    public float WristMaxStepDegrees = 1f;
    [Range(0f, 1f)]
    public float SmoothFactor = 0.5f;
    public float ProjectionEpsilon = 1e-6f;

    public SolverCCD(RobotJoint[] joints, Transform endEffector)
    {
        Joints = joints;
        EndEffector = endEffector;
    }

    public float[] SolveIK(Frame target)
    {
        int n = Joints != null ? Joints.Length : 0;
        float[] result = new float[n];

        if (Joints == null || n == 0 || EndEffector == null)
            return result;

        for (int iter = 0; iter < MaxIterations; iter++)
        {
            float posErr = Vector3.Distance(EndEffector.position, target.Position);
            if (posErr <= PositionTolerance) break;

            for (int i = n - 1; i >= 0; i--)
            {
                Transform jointT = Joints[i].JointTransform;
                Vector3 toEnd = EndEffector.position - jointT.position;
                Vector3 toTarget = target.Position - jointT.position;

                if (toEnd.sqrMagnitude < 1e-12f || toTarget.sqrMagnitude < 1e-12f) continue;

                Vector3 axisWorld = GetJointAxisWorld(Joints[i], jointT);
                float angle = SignedAngleBetweenVectorsProjected(toEnd, toTarget, axisWorld, ProjectionEpsilon);

                float maxStep = (i >= n - 3) ? WristMaxStepDegrees : MaxStepDegrees;
                float delta = Mathf.Clamp(angle, -maxStep, maxStep);

                float current = Joints[i].GetValue();
                float intended = current + delta;
                float newAngle = Mathf.Lerp(current, intended, SmoothFactor);
                newAngle = Mathf.Clamp(newAngle, Joints[i].MinLimit, Joints[i].MaxLimit);

                Joints[i].SetValue(newAngle);
            }
        }

        AlignOrientation(target);

        for (int i = 0; i < n; i++) result[i] = Joints[i].GetValue();
        return result;
    }

    void AlignOrientation(Frame target)
    {
        int n = Joints.Length;
        if (n == 0) return;

        float angError = Quaternion.Angle(EndEffector.rotation, target.Rotation);
        if (angError <= OrientationToleranceDeg) return;

        int start = Mathf.Max(0, n - 3);

        for (int iter = 0; iter < 12; iter++)
        {
            for (int i = n - 1; i >= start; i--)
            {
                Transform jointT = Joints[i].JointTransform;
                Vector3 axisWorld = GetJointAxisWorld(Joints[i], jointT);

                Quaternion fromTo = target.Rotation * Quaternion.Inverse(EndEffector.rotation);
                fromTo.ToAngleAxis(out float angleDeg, out Vector3 axis);
                if (angleDeg > 180f) angleDeg -= 360f;
                if (float.IsNaN(angleDeg) || Mathf.Abs(angleDeg) < 1e-4f) continue;

                float sign = Mathf.Sign(Vector3.Dot(axis, axisWorld));
                float step = Mathf.Clamp(angleDeg * sign, -5f, 5f);

                float maxStep = WristMaxStepDegrees;
                float delta = Mathf.Clamp(step, -maxStep, maxStep);

                float current = Joints[i].GetValue();
                float intended = current + delta;
                float newAngle = Mathf.Lerp(current, intended, SmoothFactor);
                newAngle = Mathf.Clamp(newAngle, Joints[i].MinLimit, Joints[i].MaxLimit);

                Joints[i].SetValue(newAngle);
            }

            if (Quaternion.Angle(EndEffector.rotation, target.Rotation) <= OrientationToleranceDeg) break;
        }
    }

    Vector3 GetJointAxisWorld(RobotJoint joint, Transform jointT)
    {
        switch (joint.Axis)
        {
            case RobotJoint.RotationAxis.X: return jointT.right;
            case RobotJoint.RotationAxis.Y: return jointT.up;
            case RobotJoint.RotationAxis.Z: return jointT.forward;
        }
        return jointT.up;
    }

    static float SignedAngleBetweenVectorsProjected(Vector3 v1, Vector3 v2, Vector3 axis, float eps)
    {
        Vector3 p1 = Vector3.ProjectOnPlane(v1, axis);
        Vector3 p2 = Vector3.ProjectOnPlane(v2, axis);

        if (p1.sqrMagnitude < eps || p2.sqrMagnitude < eps) return 0f;
        return Vector3.SignedAngle(p1, p2, axis.normalized);
    }
}