// 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
    {
        // --- MODIFICATION: Static variable to track the currently active attack state ---
        private static int activeAttackInstanceId = 0;
        private int myAttackInstanceId;
        // --- END MODIFICATION ---

        [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("Normalized time point to start allowing the next attack input.")]
        public float blockInputBeforeTime = 0.5f;
        
        [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;
        
        [Header("Combo Timing Window")]
        [Tooltip("Enable a special timing window at the end of the attack to chain the next combo hit.")]
        public bool useComboTimingWindow = false;

        [Tooltip("Normalized time to START the combo window.")]
        [Range(0, 1)]
        public float comboWindowStartTime = 0.8f;

        [Tooltip("How long the combo window stays open in REAL-WORLD SECONDS.")]
        public float comboWindowDuration = 0.5f;

        [Tooltip("The time scale to use during the combo window for a slow-motion effect.")]
        public float comboWindowTimeScale = 0.1f;

        [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;
        
        private bThirdPersonController _characterController;
        private bool _isRotationLockedByThis;
        private bool _comboWindowEffectTriggered;
        private Animator _animator;

        override public void OnStateEnter(Animator animator, AnimatorStateInfo stateInfo, int layerIndex)
        {
            if (_animator == null) _animator = animator;
            
            // --- MODIFICATION: Assign a unique ID to this state instance ---
            myAttackInstanceId = ++activeAttackInstanceId;
            // --- END MODIFICATION ---

            mFighter = animator.GetComponent<vIAttackListener>();
            _characterController = animator.GetComponent<bThirdPersonController>();

            if (Player.Instance != null)
                _autoTargettingInstance = Player.Instance.AutoTarget;

            if (_autoTargettingInstance == null && debug)
                Debug.LogWarning($"({damageType}) AutoTargetting instance not found. Rotation/Target features limited.");

            isAttacking = true;
            isActive = false;
            m_hasScaledTime = false;
            _comboWindowEffectTriggered = false;

            if (_characterController != null)
            {
                _characterController.lockRotation = true;
                _isRotationLockedByThis = true;
            }

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

            if (debug)
                Debug.Log($"({damageType}, ID: {myAttackInstanceId}) OnStateEnter. Now the authoritative state.");
            
            // --- MODIFICATION: Immediately block input on enter ---
            // This ensures the new state takes control right away.
            BlockAttack(true);
            // --- END MODIFICATION ---

            if (attackTimeScaleStart < 0f) attackTimeScaleStart = startDamage;
            if (attackTimeScaleEnd < 0f) attackTimeScaleEnd = endDamage;
        }

        void BlockAttack(bool block)
        {
            // --- MODIFICATION: Only allow the authoritative state to change the lock ---
            if (myAttackInstanceId != activeAttackInstanceId)
            {
                if(debug) Debug.Log($"({damageType}, ID: {myAttackInstanceId}) Tried to change block but I am not the active instance ({activeAttackInstanceId}). Ignoring.");
                return;
            }
            // --- END MODIFICATION ---

            if (Player.Instance != null)
            {
                var meleeInput = Player.Instance.MeleeCombatInput;
                if (meleeInput != null)
                {
                    if (meleeInput.BlockAttack != block) // Only log/change if there is a change
                    {
                        if(debug) Debug.Log($"({damageType}, ID: {myAttackInstanceId}) Setting BlockAttack to: {block}");
                        meleeInput.BlockAttack = block;
                    }
                }
            }
        }
        
        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;

            ManageComboLogic(currentNormalizedTime);

            if (_characterController != null && _characterController.lockRotation)
                AttemptRotationTowardsAutoTarget(animator);

            AttemptPositionLerp(animator);
            UpdateRotationLock(currentNormalizedTime);

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

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

            if (isAttacking && currentNormalizedTime > endDamage)
            {
                if (mFighter != null) mFighter.OnDisableAttack();
                isAttacking = false;
            }
        }

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

            if (isActive)
            {
                isActive = false;
                ActiveDamage(animator, false);
            }

            if (isAttacking && mFighter != null)
                mFighter.OnDisableAttack();

            isAttacking = false;
            m_hasScaledTime = false;
            
            if (_comboWindowEffectTriggered && TimeController.Instance != null)
            {
                _comboWindowEffectTriggered = false;
                TimeController.Instance.Reset();
            }

            if (_characterController != null && _isRotationLockedByThis)
            {
                _characterController.lockRotation = false;
                _isRotationLockedByThis = false;
            }

            // --- MODIFICATION: When exiting, ensure the attack is unblocked. ---
            // This is a safety net. If this was the last attack in a combo,
            // we need to make sure the input is unlocked for future actions.
            BlockAttack(false);
            // --- END MODIFICATION ---
        }
        
        private void ManageComboLogic(float currentNormalizedTime)
        {
            if (blockInputBeforeTime > 0f)
            {
                if (currentNormalizedTime >= blockInputBeforeTime)
                {
                    BlockAttack(false); // Unlock input
                }
                else
                {
                    BlockAttack(true); // Block input
                }
            }
            
            if (!useComboTimingWindow) return;

            bool isInsideWindow = currentNormalizedTime >= comboWindowStartTime;

            if (isInsideWindow)
            {
                if (!_comboWindowEffectTriggered && TimeController.Instance != null)
                {
                    _comboWindowEffectTriggered = true;
                    TimeController.Instance.SetTimeScaleForRealTimeSec(comboWindowTimeScale, comboWindowDuration, false);
                    if (debug) Debug.Log($"({damageType}) COMBO WINDOW OPEN. Accepting input.");
                }
            }
        }

        // --- Other methods remain unchanged ---

        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);
        }

        private void UpdateRotationLock(float normalizedTime)
        {
            if (_characterController != null && _isRotationLockedByThis && normalizedTime >= unlockRotationTime)
            {
                _characterController.lockRotation = false;
                _isRotationLockedByThis = false; 
            }
        }

        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);

            if (distance <= maxLerpDistance && distance > stoppingDistance)
            {
                Vector3 directionToTarget = (targetTransform.position - playerTransform.position).normalized;
                directionToTarget.y = 0;
                Vector3 targetPosition = targetTransform.position - directionToTarget * stoppingDistance;
                playerTransform.position = Vector3.MoveTowards(playerTransform.position, targetPosition, positionLerpSpeed * Time.deltaTime);
            }
        }

        private void UpdateSlowMotion(Animator animator, AnimatorStateInfo stateInfo, float currentNormalizedTime)
        {
            if (_autoTargettingInstance == null || TimeController.Instance == null) return;
            if (!m_hasScaledTime && currentNormalizedTime >= attackTimeScaleStart && currentNormalizedTime <= attackTimeScaleEnd)
            {
                bool triggerSlowMo = false;
                if (_autoTargettingInstance.CurrentTarget != null)
                {
                    float distSqr = (_autoTargettingInstance.CurrentTarget.transform.position - animator.transform.position).sqrMagnitude;
                    if (distSqr <= maxTargetDistance * maxTargetDistance)
                    {
                        triggerSlowMo = true;
                    }
                }
                if (triggerSlowMo)
                {
                    float slowMoEffectDuration = (attackTimeScaleEnd - currentNormalizedTime) * stateInfo.length;
                    if (slowMoEffectDuration > 0.01f)
                    {
                        TimeController.Instance.SetTimeScaleForSec(attackTimeScale, slowMoEffectDuration);
                        m_hasScaledTime = true;
                    }
                }
            }
            else if (m_hasScaledTime && currentNormalizedTime > attackTimeScaleEnd)
            {
                m_hasScaledTime = false; 
            }
        }

        void ActiveDamage(Animator animator, bool value)
        {
            var meleeManager = animator.GetComponent<vMeleeManager>();
            if (meleeManager)
                meleeManager.SetActiveAttack(bodyParts, meleeAttackType, value, damageMultiplier, recoilID, reactionID, ignoreDefense, activeRagdoll, senselessTime, damageType);
        }
    }
}