// Paste this code into your existing bMeleeAttackControl.cs file, replacing its content.

using System;
using System.Collections;
using System.Collections.Generic;
using Beyond; // For Player, GameStateManager, TimeController
using Invector.vCharacterController.AI.FSMBehaviour; // For vFSMBehaviourController
using UnityEngine;

namespace Invector.vMelee
{
    using vEventSystems; // For vIAttackListener

    public class bMeleeAttackControl : StateMachineBehaviour
    {
        [Header("Damage Window")]
        [Tooltip("NormalizedTime of Active Damage")]
        public float startDamage = 0.05f;
        [Tooltip("NormalizedTime of Disable Damage")]
        public float endDamage = 0.9f;

        [Header("Attack Properties")]
        public int damageMultiplier;
        public int recoilID;
        public int reactionID;
        public vAttackType meleeAttackType = vAttackType.Unarmed;
        [Tooltip("You can use a name as reference to trigger a custom HitDamageParticle")]
        public string damageType;
        [Tooltip("Body parts to activate for damage detection.")]
        public List<string> bodyParts = new List<string> { "RightLowerArm" };

        [Header("Hit Effects")]
        public bool ignoreDefense;
        public bool activeRagdoll;
        [vHideInInspector("activeRagdoll")]
        [Tooltip("Time to keep Ragdoll active if activeRagdoll is true.")]
        public float senselessTime;

        [Header("Attack Flow")]
        [Tooltip("Check true in the last attack of your combo to reset the FSM attack triggers.")]
        public bool resetAttackTrigger;
        [Tooltip("Normalized time point to reset attack triggers if resetAttackTrigger is true.")]
        public float resetTriggerBeforeTime = 0.5f;
        
        // --- NEW ---: Combo Rotation and Position Lerp Settings
        [Header("Combo & Movement")]
        [Tooltip("Normalized time to unlock rotation, allowing the player to aim the next attack in a combo. Set to 1 to disable.")]
        [Range(0,1)]
        public float unlockRotationTime = 0.7f;
        
        [Tooltip("Enable to make the character move towards the target during the attack.")]
        public bool lerpPositionTowardsTarget = false;
        
        [vHideInInspector("lerpPositionTowardsTarget")]
        [Tooltip("Max distance from the target to start moving towards it.")]
        public float maxLerpDistance = 3.5f;

        [vHideInInspector("lerpPositionTowardsTarget")]
        [Tooltip("How fast the character moves towards the target.")]
        public float positionLerpSpeed = 2.0f;
        
        [vHideInInspector("lerpPositionTowardsTarget")]
        [Tooltip("How close the character should get to the target.")]
        public float stoppingDistance = 1.2f;
        // --- END NEW ---

        [Header("Slow Motion Settings")]
        [Tooltip("Enable slow motion effect during this attack based on conditions below.")]
        public bool useAttackTimeScale = false;
        public float maxTargetDistance = 3f;
        public float lowHealthTh = 10f;
        public float attackTimeScale = 0.2f;
        public float attackTimeScaleStart = -1f;
        public float attackTimeScaleEnd = -1f;

        [Header("Rotation Settings")]
        [Tooltip("If true, the character will attempt to rotate towards the auto-target during this attack state.")]
        public bool rotatePlayerTowardsTarget;
        public float degreeThreshold = 20f;

        [Header("Debug")]
        public bool debug;

        // Private state variables
        private bool isActive;
        private vIAttackListener mFighter;
        private bool isAttacking;
        private bool m_hasScaledTime;
        private AutoTargetting _autoTargettingInstance;
        
        // --- NEW ---: Private variables for new features
        private bThirdPersonController _characterController;
        private bool _isRotationLockedByThis; // Tracks if this specific state has locked rotation
        // --- END NEW ---

        override public void OnStateEnter(Animator animator, AnimatorStateInfo stateInfo, int layerIndex)
        {
            mFighter = animator.GetComponent<vIAttackListener>();
            // --- NEW ---: Get reference to the character controller
            _characterController = animator.GetComponent<bThirdPersonController>();
            // --- END NEW ---

            if (Player.Instance != null)
            {
                _autoTargettingInstance = Player.Instance.AutoTarget;
            }
            
            if (_autoTargettingInstance == null && debug)
            {
                Debug.LogWarning($"({damageType}) AutoTargetting instance not found on {animator.name}. Rotation and target-dependent features will be limited.");
            }

            isAttacking = true;
            isActive = false;
            m_hasScaledTime = false;
            
            // --- NEW ---: Lock character rotation at the beginning of the attack
            if (_characterController != null)
            {
                _characterController.lockRotation = true;
                _isRotationLockedByThis = true;
                if (debug) Debug.Log($"({damageType}) Rotation locked by state.");
            }
            // --- END NEW ---

            if (mFighter != null)
                mFighter.OnEnableAttack();

            if (debug)
                Debug.Log($"({damageType}) OnStateEnter: {animator.name}, Layer: {layerIndex}");

            if (attackTimeScaleStart < 0f) attackTimeScaleStart = startDamage;
            if (attackTimeScaleEnd < 0f) attackTimeScaleEnd = endDamage;
        }
        
        override public void OnStateUpdate(Animator animator, AnimatorStateInfo stateInfo, int layerIndex)
        {
            if (Player.Instance.ActiveWeaponTrail)
            {
                Player.Instance.ActiveWeaponTrail.m_colorMultiplier = Color.white + Color.red * damageMultiplier;
            }
            
            float currentNormalizedTime = stateInfo.normalizedTime % 1;
            if (currentNormalizedTime == 0 && stateInfo.normalizedTime > 0.5f) currentNormalizedTime = 1f;
            
            // --- MODIFIED ---: Rotation and Movement logic is now more sophisticated
            // Only execute script-based rotation if the controller's rotation is locked
            if (_characterController != null && _characterController.lockRotation)
            {
                 AttemptRotationTowardsAutoTarget(animator);
            }
            
            // Handle position lerping
            AttemptPositionLerp(animator);

            // Handle unlocking rotation for combo aiming
            UpdateRotationLock(currentNormalizedTime);
            // --- END MODIFIED ---

            if (useAttackTimeScale)
            {
                UpdateSlowMotion(animator, stateInfo, currentNormalizedTime);
            }

            if (!isActive && currentNormalizedTime >= startDamage && currentNormalizedTime <= endDamage)
            {
                if (debug) Debug.Log($"({damageType}) Enable Damage: normTime={currentNormalizedTime:F2} (Start:{startDamage:F2}, End:{endDamage:F2})");
                isActive = true;
                ActiveDamage(animator, true);
            }
            else if (isActive && currentNormalizedTime > endDamage)
            {
                if (debug) Debug.Log($"({damageType}) Disable Damage: normTime={currentNormalizedTime:F2} > {endDamage:F2}");
                isActive = false;
                ActiveDamage(animator, false);
            }

            if (isAttacking)
            {
                if (currentNormalizedTime > endDamage) 
                {
                    if (mFighter != null) mFighter.OnDisableAttack();
                    isAttacking = false; 
                    if (debug) Debug.Log($"({damageType}) OnDisableAttack called: normTime={currentNormalizedTime:F2}");
                }
                else if (resetAttackTrigger && currentNormalizedTime >= resetTriggerBeforeTime)
                {
                    if (mFighter != null) mFighter.ResetAttackTriggers();
                    if (debug) Debug.Log($"({damageType}) ResetAttackTriggers called: normTime={currentNormalizedTime:F2}");
                }
            }
        }

        override public void OnStateExit(Animator animator, AnimatorStateInfo stateInfo, int layerIndex)
        {
            if (debug) Debug.Log($"({damageType}) OnStateExit: {animator.name}");

            if (isActive)
            {
                isActive = false;
                ActiveDamage(animator, false);
                if (debug) Debug.Log($"({damageType}) Damage disabled on StateExit.");
            }

            if (isAttacking && mFighter != null)
            {
                mFighter.OnDisableAttack();
                if (debug) Debug.Log($"({damageType}) OnDisableAttack called on StateExit (fallback).");
            }
            isAttacking = false;

            m_hasScaledTime = false;

            if (mFighter != null && resetAttackTrigger)
            {
                mFighter.ResetAttackTriggers();
                if (debug) Debug.Log($"({damageType}) ResetAttackTriggers called on StateExit due to resetAttackTrigger flag.");
            }

            // --- NEW ---: Ensure rotation is unlocked upon exiting the state
            if (_characterController != null && _isRotationLockedByThis)
            {
                _characterController.lockRotation = false;
                _isRotationLockedByThis = false;
                if (debug) Debug.Log($"({damageType}) Rotation unlocked by state on exit.");
            }
            // --- END NEW ---
        }

        private void AttemptRotationTowardsAutoTarget(Animator animator)
        {
            if (!rotatePlayerTowardsTarget || _autoTargettingInstance == null || _autoTargettingInstance.CurrentTarget == null)
            {
                return;
            }

            if (_autoTargettingInstance.IsTargetInAngle(animator.transform, _autoTargettingInstance.CurrentTarget, degreeThreshold))
            {
                _autoTargettingInstance.ExecuteRotationTowardsCurrentTarget(Time.deltaTime);
            }
        }

        // --- NEW ---
        private void UpdateRotationLock(float normalizedTime)
        {
            if (_characterController != null && _isRotationLockedByThis && normalizedTime >= unlockRotationTime)
            {
                _characterController.lockRotation = false;
                _isRotationLockedByThis = false; // Stop this state from managing the lock
                if(debug) Debug.Log($"({damageType}) Rotation unlocked for combo aiming at normTime={normalizedTime:F2}");
            }
        }

        private void AttemptPositionLerp(Animator animator)
        {
            if (!lerpPositionTowardsTarget || _characterController == null || _autoTargettingInstance == null || _autoTargettingInstance.CurrentTarget == null)
            {
                return;
            }

            Transform playerTransform = _characterController.transform;
            Transform targetTransform = _autoTargettingInstance.CurrentTarget.transform;

            float distance = Vector3.Distance(playerTransform.position, targetTransform.position);

            // Only lerp if within max distance and further than stopping distance
            if (distance <= maxLerpDistance && distance > stoppingDistance)
            {
                Vector3 directionToTarget = (targetTransform.position - playerTransform.position).normalized;
                directionToTarget.y = 0; // Keep movement on the horizontal plane

                // The target position is a point in front of the enemy at the stopping distance
                Vector3 targetPosition = targetTransform.position - directionToTarget * stoppingDistance;
                
                // Use MoveTowards for consistent speed. This adds to root motion rather than fighting it.
                playerTransform.position = Vector3.MoveTowards(playerTransform.position, targetPosition, positionLerpSpeed * Time.deltaTime);
            }
        }
        // --- END NEW ---

        private void UpdateSlowMotion(Animator animator, AnimatorStateInfo stateInfo, float currentNormalizedTime)
        {
            // ... (this method remains unchanged)
            if (_autoTargettingInstance == null || TimeController.Instance == null) return;
            if (!m_hasScaledTime)
            {
                if (currentNormalizedTime >= attackTimeScaleStart && currentNormalizedTime <= attackTimeScaleEnd)
                {
                    bool triggerSlowMo = false;
                    if (_autoTargettingInstance.CurrentTarget != null)
                    {
                        float distSqr = (_autoTargettingInstance.CurrentTarget.transform.position - animator.transform.position).sqrMagnitude;
                        bool targetNear = distSqr <= maxTargetDistance * maxTargetDistance;
                        if (targetNear)
                        {
                            float currentTargetHealth = _autoTargettingInstance.GetCurrentTargetHealth();
                            bool targetHealthLow = currentTargetHealth > 0f && currentTargetHealth < lowHealthTh;
                            if (targetHealthLow)
                            {
                                triggerSlowMo = true;
                            }
                            else 
                            {
                                triggerSlowMo = true;
                            }
                        }
                    }
                    if (triggerSlowMo)
                    {
                        float slowMoEffectDuration = (attackTimeScaleEnd - currentNormalizedTime) * stateInfo.length;
                        if (slowMoEffectDuration > 0.01f)
                        {
                            TimeController.Instance.SetTimeScaleForSec(attackTimeScale, slowMoEffectDuration);
                            if (debug) Debug.Log($"({damageType}) Slow-mo ACTIVATED. Target: {_autoTargettingInstance.CurrentTarget?.name ?? "N/A"}. Duration: {slowMoEffectDuration:F2}s. NormTime: {currentNormalizedTime:F2}");
                        }
                        else if (debug) Debug.Log($"({damageType}) Slow-mo trigger met, but calculated duration too short ({slowMoEffectDuration:F2}s). NormTime: {currentNormalizedTime:F2}");
                        m_hasScaledTime = true;
                    }
                }
            }
            else if (currentNormalizedTime > attackTimeScaleEnd && m_hasScaledTime)
            {
                m_hasScaledTime = false; 
                if (debug) Debug.Log($"({damageType}) Slow-mo window ended (normTime={currentNormalizedTime:F2}). m_hasScaledTime reset.");
            }
        }

        void ActiveDamage(Animator animator, bool value)
        {
            // ... (this method remains unchanged)
            var meleeManager = animator.GetComponent<vMeleeManager>();
            if (meleeManager)
            {
                meleeManager.SetActiveAttack(bodyParts, meleeAttackType, value, damageMultiplier, recoilID, reactionID,
                    ignoreDefense, activeRagdoll, senselessTime, damageType);
            }
            else if(debug)
            {
                Debug.LogWarning($"({damageType}) vMeleeManager not found on {animator.name}. Cannot activate/deactivate damage.");
            }
        }
    }
}