using UnityEngine;

/// <summary>
/// Two-bone IK solver for a single limb (arm or leg).
/// Works with or without an Animator.
/// Assumes a 3-bone chain: start -> mid -> end.
/// Bones are assumed to point along a configurable local primary axis,
/// with a configurable local secondary axis used as "bend up".
/// </summary>
public class OctoIkLimb : MonoBehaviour
{
    public enum Axis
    {
        XPlus,
        XMinus,
        YPlus,
        YMinus,
        ZPlus,
        ZMinus
    }

    [System.Serializable]
    private class Joints
    {
        [Tooltip("start joint of the limb (e.g. shoulder or hip).")]
        public Transform Start;

        [Tooltip("Middle joint of the limb (e.g. elbow or knee).")]
        public Transform Mid;

        [Tooltip("End joint of the limb (e.g. wrist or foot).")]
        public Transform End;
    }

    [System.Serializable]
    private class Controllers
    {
        [Tooltip("Target the end joint should try to reach.")]
        public Transform Target;

        [Tooltip("Optional pole to control the bending direction of the limb (e.g. elbow or knee direction).")]
        public Transform Pole;
    }

    [System.Serializable]
    private class Settings
    {
        [Tooltip("Local primary axis pointing toward the next bone.")]
        public Axis PrimaryAxis = Axis.YPlus;

        [Tooltip("Local secondary axis defining the bend 'up' direction.")]
        public Axis SecondaryAxis = Axis.XMinus;

        [Header("Settings")]

        [Tooltip("Global IK blend weight (0 = no IK, 1 = full IK).")]
        [Range(0f, 1f)]
        public float Weight = 1;

        [Tooltip("Apply the target rotation to the end joint.")]
        public bool OrientEnd = true;
    }

    [System.Serializable]
    private class VisualDebug
    {
        [Tooltip("Enable or disable IK debug gizmos.")]
        public bool Gizmos = true;

        [Tooltip("Scale of the axis gizmos drawn at each joint.")]
        public float Size = 0.1f;
    }

    [SerializeField] private Joints _joints;
    [SerializeField] private Controllers _controllers;
    [SerializeField] private Settings _settings;
    [SerializeField] private VisualDebug _debug;

    public Transform Target => _controllers.Target;
    public Transform Pole => _controllers.Pole;

    private float _upperLength;
    private float _lowerLength;
    private float _maxReach;

    private Quaternion _axesCorrection = Quaternion.identity;


#if UNITY_EDITOR
    void OnValidate()
    {
        RecalculateLengths();
        RecalculateAxesCorrection();
    }
#endif

    void Awake()
    {
        RecalculateLengths();
        RecalculateAxesCorrection();
    }

    void LateUpdate()
    {
        Solve();
    }

    /// <summary>
    /// Recomputes bone lengths and max reach.
    /// </summary>
    private void RecalculateLengths()
    {
        if (_joints.Start && _joints.Mid && _joints.End)
        {
            _upperLength = Vector3.Distance(_joints.Start.position, _joints.Mid.position);
            _lowerLength = Vector3.Distance(_joints.Mid.position, _joints.End.position);
            _maxReach = Mathf.Max(0f, _upperLength + _lowerLength - 0.0001f);
        }
    }

    /// <summary>
    /// Precomputes a local correction rotation so that:
    /// - local primary axis is treated as +Z (forward)
    /// - local secondary axis is treated as +Y (up)
    /// </summary>
    private void RecalculateAxesCorrection()
    {
        Vector3 pLocal = AxisToVector(_settings.PrimaryAxis).normalized;
        Vector3 sLocal = AxisToVector(_settings.SecondaryAxis).normalized;

        // Step 1: rotate primary axis to +Z
        Quaternion qPrimary = Quaternion.FromToRotation(pLocal, Vector3.forward);

        // Step 2: after this, rotate around Z so that secondary goes to +Y
        Vector3 s1 = qPrimary * sLocal;
        Vector3 s1Proj = Vector3.ProjectOnPlane(s1, Vector3.forward);

        if (s1Proj.sqrMagnitude < 0.000001f)
        {
            _axesCorrection = qPrimary;
            return;
        }

        s1Proj.Normalize();
        Quaternion qTwist = Quaternion.FromToRotation(s1Proj, Vector3.up);

        // This maps (primary, secondary) -> (Z+, Y+)
        _axesCorrection = qTwist * qPrimary;
    }

    /// <summary>
    /// Solves IK for the current limb configuration.
    /// Can be called manually if autoSolve is disabled.
    /// </summary>
    public void Solve()
    {
        if (_settings.Weight <= 0f)
            return;

        if (!_joints.Start || !_joints.Mid || !_joints.End || !_controllers.Target)
            return;

        if (_upperLength <= 0f || _lowerLength <= 0f)
            RecalculateLengths();

        Vector3 startPos = _joints.Start.position;
        Vector3 midPos = _joints.Mid.position;
        Vector3 endPos = _joints.End.position;

        Vector3 startToTarget = _controllers.Target.position - startPos;
        float dist = startToTarget.magnitude;
        if (dist < 0.0001f)
            return;

        dist = Mathf.Min(dist, _maxReach);
        Vector3 dirToTarget = startToTarget.normalized;
        
        Vector3 bendNormal;

        if (_controllers.Pole)
        {
            Vector3 startToPole = (_controllers.Pole.position - startPos);
            Vector3 side = Vector3.ProjectOnPlane(startToPole, dirToTarget);
            if (side.sqrMagnitude < 0.000001f)
            {
                // Fallback if pole is collinear with forward
                side = Vector3.right;
            }
            side.Normalize();

            // Bend normal is perpendicular to forward and to that side
            bendNormal = Vector3.Cross(dirToTarget, side);
        }
        else
        {
            // Fallback: use start.up to define a side, then build a normal
            Vector3 side = Vector3.ProjectOnPlane(_joints.Start.up, dirToTarget);
            if (side.sqrMagnitude < 0.000001f)
                side = Vector3.right;
            side.Normalize();

            bendNormal = Vector3.Cross(dirToTarget, side);
        }

        if (bendNormal.sqrMagnitude < 0.000001f)
            return;

        bendNormal.Normalize();

        // Bend direction lies in the plane defined by dirToTarget and bendNormal
        Vector3 bendDir = Vector3.Cross(bendNormal, dirToTarget).normalized;

        float a = _upperLength;
        float b = _lowerLength;
        float c = dist;

        // Law of cosines to compute start angle
        float cos0 = Mathf.Clamp((a * a + c * c - b * b) / (2f * a * c), -1f, 1f);
        float angle0 = Mathf.Acos(cos0);

        float along = Mathf.Cos(angle0) * a;
        float height = Mathf.Sin(angle0) * a;

        Vector3 desiredMid = startPos + dirToTarget * along + bendDir * height;
        Vector3 desiredEnd = _controllers.Target.position;

        // Canonical solver rotations assuming +Z forward, +Y up
        Quaternion startVirtual = Quaternion.LookRotation(desiredMid - startPos, bendNormal);
        Quaternion midVirtual = Quaternion.LookRotation(desiredEnd - desiredMid, bendNormal);

        // Map canonical frame to bone-local frame using axesCorrection
        // (we expect bones to point along primaryAxis and bend using secondaryAxis)
        Quaternion startDesired = startVirtual * _axesCorrection;
        Quaternion midDesired = midVirtual * _axesCorrection;

        _joints.Start.rotation = Quaternion.Slerp(_joints.Start.rotation, startDesired, _settings.Weight);
        _joints.Mid.rotation = Quaternion.Slerp(_joints.Mid.rotation, midDesired, _settings.Weight);

        if (_settings.OrientEnd)
        {
            _joints.End.rotation = Quaternion.Slerp(_joints.End.rotation, _controllers.Target.rotation, _settings.Weight);
        }
    }

    /// <summary>
    /// Converts a LocalAxis6 enum to a unit Vector3.
    /// </summary>
    private Vector3 AxisToVector(Axis axis)
    {
        switch (axis)
        {
            default: return Vector3.forward;
            case Axis.XPlus: return Vector3.right;
            case Axis.XMinus: return Vector3.left;
            case Axis.YPlus: return Vector3.up;
            case Axis.YMinus: return Vector3.down;
            case Axis.ZMinus: return Vector3.back;
        }
    }

    void OnDrawGizmos()
    {
        if (!_debug.Gizmos)
            return;

        if (!_joints.Start || !_joints.Mid || !_joints.End)
            return;

        // Bones to Controllers
        Gizmos.color = Color.yellow;
        if (_controllers.Target)
            Gizmos.DrawLine(_joints.End.position, _controllers.Target.position);
        if (_controllers.Pole)
            Gizmos.DrawLine(_joints.Mid.position, _controllers.Pole.position);

        // Bones
        Gizmos.color = Color.white;
        Gizmos.DrawLine(_joints.Start.position, _joints.Mid.position);
        Gizmos.DrawLine(_joints.Mid.position, _joints.End.position);
        Gizmos.matrix = _joints.Start.localToWorldMatrix;
        Gizmos.DrawWireSphere(Vector3.zero, _debug.Size * .15f);
        Gizmos.matrix = _joints.Mid.localToWorldMatrix;
        Gizmos.DrawWireSphere(Vector3.zero, _debug.Size * .15f);
        Gizmos.matrix = _joints.End.localToWorldMatrix;
        Gizmos.DrawWireSphere(Vector3.zero, _debug.Size * .15f);

        // Target
        if (_controllers.Target)
        {
            Gizmos.color = Color.yellow;
            Gizmos.matrix = _controllers.Target.localToWorldMatrix;
            Gizmos.DrawWireCube(Vector3.zero, Vector3.one * _debug.Size);
        }

        // Pole
        if (_controllers.Pole)
        {
            Gizmos.color = Color.yellow;
            Gizmos.matrix = _controllers.Pole.localToWorldMatrix;
            Gizmos.DrawWireCube(Vector3.zero, Vector3.one * _debug.Size * .5f);
        }
    }
}