using UnityEngine;

public enum ArmState
{
    Environmental,
    TransitionToEquipped,
    Equipped,
    TransitionToEnvironmental
}

[System.Serializable]
public class ArmEquipment
{
    [Header("Arm References")]
    public OctoIkLimb IkLimb;
    public Transform HandSocket;
    
    [Header("Bone References (for procedural animation)")]
    [Tooltip("Shoulder bone (optional)")]
    public Transform ShoulderBone;
    
    [Tooltip("Elbow bone (optional)")]
    public Transform ElbowBone;
    
    [Tooltip("Wrist bone (required for rotation)")]
    public Transform WristBone;
    
    [Header("Equipped Pose Settings")]
    [Tooltip("Rotation to apply to shoulder when equipped")]
    public Vector3 ShoulderRotationOffset = Vector3.zero;
    
    [Tooltip("Rotation to apply to elbow when equipped")]
    public Vector3 ElbowRotationOffset = Vector3.zero;
    
    [Tooltip("Rotation to apply to wrist when equipped")]
    public Vector3 WristRotationOffset = new Vector3(0, 0, 90);
    
    [Tooltip("Speed of rotation lerp")]
    [Range(1f, 20f)]
    public float RotationSpeed = 10f;
    
    [Tooltip("Duration of IK blend after unequip transition")]
    [Range(0.1f, 1f)]
    public float IkBlendDuration = 0.3f;
    
    [Header("Transition Settings")]
    [Tooltip("Position offset where arm exits camera view")]
    public Vector3 TransitionOffset = new Vector3(0.5f, -0.3f, 0.2f);
    
    [Tooltip("Duration for arm to move out of view")]
    public float TransitionOutDuration = 0.3f;
    
    [Tooltip("Duration for arm to return to view")]
    public float TransitionInDuration = 0.3f;
    
    [Tooltip("Curve for transition movement")]
    public AnimationCurve TransitionCurve = AnimationCurve.EaseInOut(0, 0, 1, 1);
    
    [HideInInspector] public ArmState State = ArmState.Environmental;
    [HideInInspector] public Vector3 OriginalTargetPosition;
    [HideInInspector] public Quaternion OriginalTargetRotation;
    [HideInInspector] public GameObject EquippedObject;
    [HideInInspector] public float IkWeight = 1f;
    
    [HideInInspector] public Quaternion OriginalShoulderRotation;
    [HideInInspector] public Quaternion OriginalElbowRotation;
    [HideInInspector] public Quaternion OriginalWristRotation;
    
    private int _transitionTweenId = -1;
    private int _ikWeightTweenId = -1;
    private int _ikBlendTweenId = -1;
    private bool _isBlendingToIk = false;
    
    private int _tapTweenId = -1;
    private bool _isShaking = false;
    private TapAnimationSettings _currentShakeSettings;
    
    private Quaternion _transitionStartShoulder;
    private Quaternion _transitionStartElbow;
    private Quaternion _transitionStartWrist;
    
    public void Initialize()
    {
        if (IkLimb != null && IkLimb.Target != null)
        {
            OriginalTargetPosition = IkLimb.Target.localPosition;
            OriginalTargetRotation = IkLimb.Target.localRotation;
        }
        
        if (ShoulderBone != null)
            OriginalShoulderRotation = ShoulderBone.localRotation;
        
        if (ElbowBone != null)
            OriginalElbowRotation = ElbowBone.localRotation;
        
        if (WristBone != null)
            OriginalWristRotation = WristBone.localRotation;
        
        State = ArmState.Environmental;
        IkWeight = 1f;
    }
    
    public void StartTransitionToEquipped(GameObject objectToEquip, System.Action onTransitionComplete)
    {
        if (State != ArmState.Environmental)
            return;
        
        State = ArmState.TransitionToEquipped;
        EquippedObject = objectToEquip;
        
        CancelActiveTweens();
        
        if (IkLimb == null || IkLimb.Target == null)
        {
            CompleteEquip(onTransitionComplete);
            return;
        }
        
        Vector3 transitionTarget = OriginalTargetPosition + TransitionOffset;
        
        _transitionTweenId = LeanTween.value(IkLimb.gameObject, 0f, 1f, TransitionOutDuration)
            .setEase(TransitionCurve)
            .setOnUpdate((float t) => {
                if (IkLimb != null && IkLimb.Target != null)
                {
                    IkLimb.Target.localPosition = Vector3.Lerp(IkLimb.Target.localPosition, transitionTarget, t);
                }
            })
            .setOnComplete(() => {
                _transitionTweenId = -1;
                CompleteEquip(onTransitionComplete);
            }).id;
        
        _ikWeightTweenId = LeanTween.value(IkLimb.gameObject, IkWeight, 0f, TransitionOutDuration * 0.5f)
            .setOnUpdate((float val) => {
                IkWeight = val;
            })
            .setOnComplete(() => {
                _ikWeightTweenId = -1;
            }).id;
    }
    
    private void CompleteEquip(System.Action onComplete)
    {
        State = ArmState.Equipped;
        
        if (IkLimb != null)
        {
            IkLimb.enabled = false;
        }
        
        if (EquippedObject != null && HandSocket != null)
        {
            EquippedObject.transform.SetParent(HandSocket, false);
            EquippedObject.transform.localPosition = Vector3.zero;
            EquippedObject.transform.localRotation = Quaternion.identity;
        }
        
        onComplete?.Invoke();
    }
    
    public void UpdateEquippedPose()
    {
        if (State != ArmState.Equipped && !_isBlendingToIk)
            return;
        
        if (_isShaking)
        {
            UpdateShake();
            return;
        }
        
        if (State == ArmState.Equipped)
        {
            float deltaTime = UnityEngine.Time.deltaTime;
            float lerpSpeed = 1f - Mathf.Exp(-RotationSpeed * deltaTime);
            
            if (ShoulderBone != null)
            {
                Quaternion targetRot = OriginalShoulderRotation * Quaternion.Euler(ShoulderRotationOffset);
                ShoulderBone.localRotation = Quaternion.Slerp(ShoulderBone.localRotation, targetRot, lerpSpeed);
            }
            
            if (ElbowBone != null)
            {
                Quaternion targetRot = OriginalElbowRotation * Quaternion.Euler(ElbowRotationOffset);
                ElbowBone.localRotation = Quaternion.Slerp(ElbowBone.localRotation, targetRot, lerpSpeed);
            }
            
            if (WristBone != null)
            {
                Quaternion targetRot = OriginalWristRotation * Quaternion.Euler(WristRotationOffset);
                WristBone.localRotation = Quaternion.Slerp(WristBone.localRotation, targetRot, lerpSpeed);
            }
        }
    }
    
    public void StartTransitionToEnvironmental(System.Action onTransitionComplete)
    {
        if (State != ArmState.Equipped)
            return;
        
        State = ArmState.TransitionToEnvironmental;
        
        CancelActiveTweens();
        
        _transitionStartShoulder = ShoulderBone != null ? ShoulderBone.localRotation : Quaternion.identity;
        _transitionStartElbow = ElbowBone != null ? ElbowBone.localRotation : Quaternion.identity;
        _transitionStartWrist = WristBone != null ? WristBone.localRotation : Quaternion.identity;
        
        if (IkLimb == null || IkLimb.Target == null)
        {
            CompleteUnequip(onTransitionComplete);
            return;
        }
        
        Vector3 currentPosition = IkLimb.Target.localPosition;
        Vector3 transitionTarget = OriginalTargetPosition + TransitionOffset;
        
        _ikWeightTweenId = LeanTween.value(IkLimb.gameObject, IkWeight, 0f, TransitionInDuration * 0.5f)
            .setOnUpdate((float val) => {
                IkWeight = val;
            })
            .setOnComplete(() => {
                _ikWeightTweenId = -1;
            }).id;
        
        _transitionTweenId = LeanTween.value(IkLimb.gameObject, 0f, 1f, TransitionInDuration)
            .setEase(TransitionCurve)
            .setOnUpdate((float t) => {
                if (IkLimb != null && IkLimb.Target != null)
                {
                    IkLimb.Target.localPosition = Vector3.Lerp(currentPosition, transitionTarget, t);
                }
                
                if (ShoulderBone != null)
                    ShoulderBone.localRotation = Quaternion.Slerp(_transitionStartShoulder, OriginalShoulderRotation, t);
                
                if (ElbowBone != null)
                    ElbowBone.localRotation = Quaternion.Slerp(_transitionStartElbow, OriginalElbowRotation, t);
                
                if (WristBone != null)
                    WristBone.localRotation = Quaternion.Slerp(_transitionStartWrist, OriginalWristRotation, t);
            })
            .setOnComplete(() => {
                _transitionTweenId = -1;
                StartIkBlend(onTransitionComplete);
            }).id;
    }
    
    private void StartIkBlend(System.Action onTransitionComplete)
    {
        if (IkLimb != null)
        {
            IkLimb.enabled = true;
        }
        
        if (EquippedObject != null)
        {
            if (HandSocket != null)
                EquippedObject.transform.SetParent(null);
            
            UnityEngine.Object.Destroy(EquippedObject);
            EquippedObject = null;
        }
        
        _isBlendingToIk = true;
        
        Quaternion startShoulder = ShoulderBone != null ? ShoulderBone.localRotation : Quaternion.identity;
        Quaternion startElbow = ElbowBone != null ? ElbowBone.localRotation : Quaternion.identity;
        Quaternion startWrist = WristBone != null ? WristBone.localRotation : Quaternion.identity;
        Vector3 startTargetPos = IkLimb != null && IkLimb.Target != null ? IkLimb.Target.localPosition : Vector3.zero;
        Quaternion startTargetRot = IkLimb != null && IkLimb.Target != null ? IkLimb.Target.localRotation : Quaternion.identity;
        
        if (IkLimb != null && IkLimb.gameObject != null)
        {
            _ikBlendTweenId = LeanTween.value(IkLimb.gameObject, 0f, 1f, IkBlendDuration)
                .setEase(LeanTweenType.easeOutQuad)
                .setOnUpdate((float t) => {
                    if (ShoulderBone != null)
                        ShoulderBone.localRotation = Quaternion.Slerp(startShoulder, OriginalShoulderRotation, t);
                    
                    if (ElbowBone != null)
                        ElbowBone.localRotation = Quaternion.Slerp(startElbow, OriginalElbowRotation, t);
                    
                    if (WristBone != null)
                        WristBone.localRotation = Quaternion.Slerp(startWrist, OriginalWristRotation, t);
                    
                    if (IkLimb != null && IkLimb.Target != null)
                    {
                        IkLimb.Target.localPosition = Vector3.Lerp(startTargetPos, OriginalTargetPosition, t);
                        IkLimb.Target.localRotation = Quaternion.Slerp(startTargetRot, OriginalTargetRotation, t);
                    }
                    
                    IkWeight = t;
                })
                .setOnComplete(() => {
                    _ikBlendTweenId = -1;
                    CompleteUnequip(onTransitionComplete);
                }).id;
        }
        else
        {
            CompleteUnequip(onTransitionComplete);
        }
    }
    
    private void CompleteUnequip(System.Action onComplete)
    {
        State = ArmState.Environmental;
        _isBlendingToIk = false;
        
        if (IkLimb != null && IkLimb.Target != null)
        {
            IkLimb.Target.localPosition = OriginalTargetPosition;
            IkLimb.Target.localRotation = OriginalTargetRotation;
        }
        
        if (ShoulderBone != null)
            ShoulderBone.localRotation = OriginalShoulderRotation;
        
        if (ElbowBone != null)
            ElbowBone.localRotation = OriginalElbowRotation;
        
        if (WristBone != null)
            WristBone.localRotation = OriginalWristRotation;
        
        onComplete?.Invoke();
    }
    
    public void CancelActiveTweens()
    {
        if (_transitionTweenId != -1)
        {
            LeanTween.cancel(_transitionTweenId);
            _transitionTweenId = -1;
        }
        
        if (_ikWeightTweenId != -1)
        {
            LeanTween.cancel(_ikWeightTweenId);
            _ikWeightTweenId = -1;
        }
        
        if (_ikBlendTweenId != -1)
        {
            LeanTween.cancel(_ikBlendTweenId);
            _ikBlendTweenId = -1;
        }
        
        if (_tapTweenId != -1)
        {
            LeanTween.cancel(_tapTweenId);
            _tapTweenId = -1;
        }
        
        _isBlendingToIk = false;
        _isShaking = false;
    }
    
    public void PlayTapAnimation(TapAnimationSettings settings)
    {
        if (State != ArmState.Equipped || settings == null)
            return;
        
        if (_tapTweenId != -1)
        {
            LeanTween.cancel(_tapTweenId);
        }
        
        Quaternion elbowStart = ElbowBone != null ? ElbowBone.localRotation : Quaternion.identity;
        Quaternion wristStart = WristBone != null ? WristBone.localRotation : Quaternion.identity;
        
        Quaternion elbowWindup = elbowStart * Quaternion.Euler(0, 0, -settings.WindupRotation);
        Quaternion elbowStrike = elbowStart * Quaternion.Euler(0, 0, settings.WindupRotation * 0.3f);
        
        Quaternion wristWindup = wristStart;
        Quaternion wristStrike = wristStart * Quaternion.Euler(0, 0, -settings.WristSnapRotation);
        
        _tapTweenId = LeanTween.value(IkLimb.gameObject, 0f, 1f, settings.WindupDuration)
            .setOnUpdate((float t) => {
                if (ElbowBone != null)
                {
                    ElbowBone.localRotation = Quaternion.Slerp(elbowStart, elbowWindup, t);
                }
                if (WristBone != null)
                {
                    WristBone.localRotation = Quaternion.Slerp(wristStart, wristWindup, t);
                }
            })
            .setOnComplete(() => {
                _tapTweenId = LeanTween.value(IkLimb.gameObject, 0f, 1f, settings.StrikeDuration)
                    .setOnUpdate((float t) => {
                        if (ElbowBone != null)
                        {
                            ElbowBone.localRotation = Quaternion.Slerp(elbowWindup, elbowStrike, t);
                        }
                        if (WristBone != null)
                        {
                            WristBone.localRotation = Quaternion.Slerp(wristWindup, wristStrike, t);
                        }
                    })
                    .setOnComplete(() => {
                        _tapTweenId = LeanTween.value(IkLimb.gameObject, 0f, 1f, 0.1f)
                            .setOnUpdate((float t) => {
                                if (ElbowBone != null)
                                {
                                    ElbowBone.localRotation = Quaternion.Slerp(elbowStrike, elbowStart, t);
                                }
                                if (WristBone != null)
                                {
                                    WristBone.localRotation = Quaternion.Slerp(wristStrike, wristStart, t);
                                }
                            })
                            .setOnComplete(() => {
                                _tapTweenId = -1;
                            }).id;
                    }).id;
            }).id;
    }
    
    public void StartShake(TapAnimationSettings settings)
    {
        _isShaking = true;
        _currentShakeSettings = settings;
    }
    
    public void StopShake()
    {
        _isShaking = false;
        _currentShakeSettings = null;
    }
    
    public void UpdateShake()
    {
        if (!_isShaking || State != ArmState.Equipped || WristBone == null || _currentShakeSettings == null)
            return;
        
        float shake = Mathf.Sin(UnityEngine.Time.time * _currentShakeSettings.ShakeSpeed) * _currentShakeSettings.ShakeIntensity;
        Quaternion shakeRotation = Quaternion.Euler(shake, shake * 0.5f, 0);
        
        Quaternion targetRot = OriginalWristRotation * Quaternion.Euler(WristRotationOffset) * shakeRotation;
        WristBone.localRotation = targetRot;
    }
}