using UnityEngine;
using System.Linq;

namespace KindRetargeting
{
    [DefaultExecutionOrder(3)]
    public class FullBodyInverseKinematics : MonoBehaviour
    {
        [System.Serializable]
        public class IKTarget
        {
            public Transform target;
            [Range(0, 1)] public float weight = 1f;
        }

        [System.Serializable]
        public class LimbIK
        {
            public bool enableIK = true;  // IK 활성화 여부
            public IKTarget endTarget;     // 손/발 타겟
            public IKTarget middleTarget;  // 팔꿈치/무릎 타겟
            public IKTarget ankleTarget;   // 발목 타겟 (다리용)
        }

        [System.Serializable]
        public class FootIK : LimbIK
        {
            [Range(0, 1)] public float footRotationWeight = 1f;  // 발 회전 가중치만 남김
        }

        [HideInInspector] public Animator animator;

        [Header("팔 타겟")]
        public LimbIK leftArm = new LimbIK();
        public LimbIK rightArm = new LimbIK();

        [Header("다리 타겟")]
        public LimbIK leftLeg = new LimbIK();
        public LimbIK rightLeg = new LimbIK();

        [Header("IK 설정")]
        [Range(1, 20)] public int iterations = 10;
        public float deltaThreshold = 0.001f;

        private Transform[][] limbs;
        private Vector3[][] positions;
        private float[][] boneLengths;
        private Vector3[][] initialDirections;
        private Quaternion[][] initialRotations;

        private void Start()
        {
            if (!ValidateAnimator()) return;

            animator = GetComponent<Animator>();

            // 팔과 다리의 본 체인 초기화
            limbs = new Transform[][]
            {
                GetBoneChain(HumanBodyBones.LeftUpperArm, HumanBodyBones.LeftLowerArm, HumanBodyBones.LeftHand),
                GetBoneChain(HumanBodyBones.RightUpperArm, HumanBodyBones.RightLowerArm, HumanBodyBones.RightHand),
                GetBoneChain(HumanBodyBones.LeftUpperLeg, HumanBodyBones.LeftLowerLeg, HumanBodyBones.LeftFoot, HumanBodyBones.LeftToes),
                GetBoneChain(HumanBodyBones.RightUpperLeg, HumanBodyBones.RightLowerLeg, HumanBodyBones.RightFoot, HumanBodyBones.RightToes)
            };

            InitializeArrays();
        }

        private bool ValidateAnimator()
        {
            if (animator == null || !animator.avatar || !animator.avatar.isHuman)
            {
                Debug.LogError("유효한 Humanoid Animator가 필요합니다.");
                enabled = false;
                return false;
            }
            return true;
        }

        private Transform[] GetBoneChain(params HumanBodyBones[] bones)
        {
            return bones.Select(bone => animator.GetBoneTransform(bone)).ToArray();
        }

        private void InitializeArrays()
        {
            positions = new Vector3[limbs.Length][];
            boneLengths = new float[limbs.Length][];
            initialDirections = new Vector3[limbs.Length][];
            initialRotations = new Quaternion[limbs.Length][];

            for (int j = 0; j < limbs.Length; j++)
            {
                int chainLength = limbs[j].Length;
                positions[j] = new Vector3[chainLength];
                boneLengths[j] = new float[chainLength - 1];
                initialDirections[j] = new Vector3[chainLength];
                initialRotations[j] = new Quaternion[chainLength];

                // 본 길이와 초기 방향 저장
                for (int i = 0; i < chainLength - 1; i++)
                {
                    Vector3 boneVector = limbs[j][i + 1].position - limbs[j][i].position;
                    Vector3 animatorScale = animator.transform.lossyScale;
                    float averageScale = (animatorScale.x + animatorScale.y + animatorScale.z) / 3f;
                    boneLengths[j][i] = boneVector.magnitude / averageScale;
                    initialDirections[j][i] = limbs[j][i].InverseTransformDirection(
                        (limbs[j][i + 1].position - limbs[j][i].position).normalized
                    );
                    initialRotations[j][i] = limbs[j][i].rotation;
                }
                initialRotations[j][chainLength - 1] = limbs[j][chainLength - 1].rotation;
            }
        }

        private void Update()
        {
            if (!ValidateAnimator()) return;
            
            // IK 계산 및 적용
            SolveIK();
        }

        private void SolveIK()
        {
            if (animator == null) return;

            LimbIK[] limbTargets = { leftArm, rightArm, leftLeg, rightLeg };

            for (int j = 0; j < limbs.Length; j++)
            {
                // IK가 비활성화되어 있으면 해당 팔다리는 처리하지 않고 스킵
                if (!limbTargets[j].enableIK) continue;

                // 타겟이 없으면 스킵
                if (limbTargets[j].endTarget.target == null) continue;

                // 현재 위치 저장
                for (int i = 0; i < limbs[j].Length; i++)
                    positions[j][i] = limbs[j][i].position;

                bool isArm = j < 2;
                SolveIK(j, limbTargets[j], isArm);
                ApplyRotations(j, limbTargets[j]);
            }
        }

        private void SolveIK(int chainIndex, LimbIK limbIK, bool isArm)
        {
            // 현재 애니메이터 스케일 계산
            Vector3 animatorScale = animator.transform.lossyScale;
            float averageScale = (animatorScale.x + animatorScale.y + animatorScale.z) / 3f;

            // 중간 관절 처리
            if (limbIK.middleTarget.target != null && limbIK.middleTarget.weight > 0)
            {
                if (isArm)
                    ProcessArmElbow(chainIndex, limbIK.middleTarget);
                else
                    ProcessLegKnee(chainIndex, limbIK.middleTarget);
            }

            // 발목 타겟 처리 (다리인 경우에만)
            if (!isArm && limbIK is FootIK footIK && footIK.ankleTarget.target != null && footIK.ankleTarget.weight > 0)
            {
                ProcessAnkle(chainIndex, footIK.ankleTarget);
            }

            // FABRIK 반복
            for (int iteration = 0; iteration < iterations; iteration++)
            {
                // Forward
                positions[chainIndex][positions[chainIndex].Length - 1] = Vector3.Lerp(
                    positions[chainIndex][positions[chainIndex].Length - 1],
                    limbIK.endTarget.target.position,
                    limbIK.endTarget.weight
                );

                for (int i = positions[chainIndex].Length - 2; i >= 0; i--)
                {
                    Vector3 direction = (positions[chainIndex][i] - positions[chainIndex][i + 1]).normalized;
                    positions[chainIndex][i] = positions[chainIndex][i + 1] + direction * (boneLengths[chainIndex][i] * averageScale);
                }

                // Backward
                positions[chainIndex][0] = limbs[chainIndex][0].position;
                for (int i = 1; i < positions[chainIndex].Length; i++)
                {
                    Vector3 direction = (positions[chainIndex][i] - positions[chainIndex][i - 1]).normalized;
                    positions[chainIndex][i] = positions[chainIndex][i - 1] + direction * (boneLengths[chainIndex][i - 1] * averageScale);
                }

                if ((positions[chainIndex][positions[chainIndex].Length - 1] - limbIK.endTarget.target.position).sqrMagnitude
                    < deltaThreshold * deltaThreshold)
                    break;
            }
        }

        private void ProcessArmElbow(int chainIndex, IKTarget middleTarget)
        {
            if (middleTarget?.target == null || middleTarget.weight <= 0) return;

            // 팔 상완(Upper Arm)을 기준으로 계산
            Vector3 upperArmPos = positions[chainIndex][0];  // 상완 시작점 (이전의 1에서 0으로 변경)
            Vector3 elbowPos = positions[chainIndex][1];     // 팔꿈치 위치 (이전의 2에서 1로 변경)
            Vector3 handPos = positions[chainIndex][2];      // 손 위치 (이전의 3에서 2로 변경)

            // 팔꿈치 타겟 방향 계산
            Vector3 currentElbowDir = (elbowPos - upperArmPos).normalized;
            Vector3 targetElbowDir = (middleTarget.target.position - upperArmPos).normalized;

            // 팔꿈치 회전 계산
            Quaternion elbowRotation = Quaternion.FromToRotation(currentElbowDir, targetElbowDir);
            float elbowLength = Vector3.Distance(upperArmPos, elbowPos);

            // 새로운 팔꿈치 위치 계산
            Vector3 newElbowPos = upperArmPos + (elbowRotation * currentElbowDir) * elbowLength;

            // 가중치를 적용하여 팔꿈치 위치 보간
            positions[chainIndex][1] = Vector3.Lerp(elbowPos, newElbowPos, middleTarget.weight);
        }

        private void ProcessLegKnee(int chainIndex, IKTarget middleTarget)
        {
            Vector3 rootPos = positions[chainIndex][0];
            Vector3 middlePos = positions[chainIndex][1];

            Vector3 currentDirection = (middlePos - rootPos).normalized;
            Vector3 targetDirection = (middleTarget.target.position - rootPos).normalized;

            Quaternion middleRotation = Quaternion.FromToRotation(currentDirection, targetDirection);
            positions[chainIndex][1] = rootPos +
                Vector3.Lerp(currentDirection, middleRotation * currentDirection, middleTarget.weight) *
                Vector3.Distance(rootPos, middlePos);
        }

        private void ProcessAnkle(int chainIndex, IKTarget ankleTarget)
        {
            if (ankleTarget?.target == null || ankleTarget.weight <= 0) return;

            Vector3 anklePos = positions[chainIndex][2];  // 발목 위치
            Vector3 newAnklePos = Vector3.Lerp(anklePos, ankleTarget.target.position, ankleTarget.weight);
            positions[chainIndex][2] = newAnklePos;
        }

        private void ApplyRotations(int chainIndex, LimbIK limbIK)
        {
            // 기본 회전 로직만 남김
            for (int i = 0; i < limbs[chainIndex].Length - 1; i++)
            {
                ApplyBasicRotation(chainIndex, i);
            }

            // 마지막 본(손/발)의 회전
            var lastBone = limbs[chainIndex].Length - 1;
            limbs[chainIndex][lastBone].position = positions[chainIndex][lastBone];
            limbs[chainIndex][lastBone].rotation = Quaternion.Lerp(
                limbs[chainIndex][lastBone].rotation,
                limbIK.endTarget.target.rotation,
                limbIK.endTarget.weight
            );
        }

        private void ApplyBasicRotation(int chainIndex, int boneIndex)
        {
            limbs[chainIndex][boneIndex].position = positions[chainIndex][boneIndex];

            Vector3 toNext = (positions[chainIndex][boneIndex + 1] - positions[chainIndex][boneIndex]).normalized;
            Quaternion fromToRotation = Quaternion.FromToRotation(
                limbs[chainIndex][boneIndex].TransformDirection(initialDirections[chainIndex][boneIndex]),
                toNext
            );

            limbs[chainIndex][boneIndex].rotation = fromToRotation * limbs[chainIndex][boneIndex].rotation;
        }
    }
}
