DeviantMobile-Rohan/Assets/Scripts/Weapons/WeaponManager.cs

using System.Collections.Generic;
using UnityEngine;
using UnityEngine.InputSystem;

/// <summary>
/// Manages equipping/unequipping weapons on the player and overrides attack animations
/// so that only the equipped weapon's animations are used.
/// Assumes weapon models (bat/hammer) are already attached to the player's hand and hidden.
/// </summary>
public class WeaponManager : MonoBehaviour
{
    public enum WeaponType { None, Bat, Hammer }

    [Header("In-hand weapon objects (already attached to player hand)")]
    [SerializeField] private GameObject batInHand;
    [SerializeField] private GameObject hammerInHand;

    [Header("Optional: disable base attack inputs while a weapon is equipped")]
    [SerializeField] private bool blockBaseAttacksWhenEquipped = true;

    [Header("Input System")]
    [SerializeField] private string pickupActionName = "WeaponPickup";

    [Header("Pickup/Drop Settings")]
    [SerializeField] private float dropReequipBlockSeconds = 0.2f;

    // Tracks the current equipped weapon
    public WeaponType CurrentWeapon { get; private set; } = WeaponType.None;
    public bool IsWeaponEquipped => CurrentWeapon != WeaponType.None;

    // Internal state for alternating swings (e.g., BatOne/BatTwo)
    private int swingIndex = 0;

    // Base attacks to override when a weapon is equipped
    private static readonly HashSet<string> baseAttackNames = new HashSet<string>
    {
        AnimationNames.Jab,
        AnimationNames.Haymaker,
        AnimationNames.BasicKick,
        AnimationNames.SumoSlap,
        AnimationNames.SpinningBackfist,
        AnimationNames.ElbowSmash,
        AnimationNames.AirPunch,
        AnimationNames.LowKick,
        // common fallbacks
        "Punch", "Kick", "Strike"
    };

    private PlayerInput _playerInput;
    private InputAction _pickupAction;
    private float _pickupCooldownUntil = 0f;

    private void Awake()
    {
        _playerInput = GetComponent<PlayerInput>();
        if (_playerInput != null)
        {
            _pickupAction = _playerInput.actions?.FindAction(pickupActionName, false);
            if (_pickupAction != null && !_pickupAction.enabled)
                _pickupAction.SafeEnable();
        }
    }

    private void OnEnable()
    {
        if (_pickupAction != null)
            _pickupAction.performed += OnPickupOrDrop;
    }

    private void OnDisable()
    {
        if (_pickupAction != null)
            _pickupAction.performed -= OnPickupOrDrop;
    }

    private void Start()
    {
        // Ensure correct visibility at start
        ApplyVisibility();
    }

    /// <summary>
    /// Equip a specific weapon and reveal it on the player's hand.
    /// </summary>
    public void EquipWeapon(WeaponType weapon)
    {
        if (CurrentWeapon == weapon) return;

        CurrentWeapon = weapon;
        swingIndex = 0; // reset combo index on equip
        ApplyVisibility();
        Debug.Log($"Equipped weapon: {CurrentWeapon}");
    }

    /// <summary>
    /// Unequip any weapon and restore default (unarmed) state.
    /// </summary>
    public void UnequipWeapon()
    {
        if (CurrentWeapon == WeaponType.None) return;
        CurrentWeapon = WeaponType.None;
        swingIndex = 0;
        ApplyVisibility();
        Debug.Log("Unequipped weapon");
    }

    /// <summary>
    /// Resolves an incoming requested animation to a weapon-specific animation if a weapon is equipped.
    /// Non-attack animations are passed through unchanged.
    /// </summary>
    public string ResolveAnimation(string requested)
    {
        if (!IsWeaponEquipped)
            return requested;

        // Allow reactions, blocks, locomotion and non-attack anims
        if (!baseAttackNames.Contains(requested))
            return requested;

        switch (CurrentWeapon)
        {
            case WeaponType.Bat:
                return NextBatAnimation();
            case WeaponType.Hammer:
                return NextHammerAnimation();
            default:
                return requested;
        }
    }

    private string NextBatAnimation()
    {
        // Alternate between two bat swings
        swingIndex = (swingIndex + 1) % 2;
        return swingIndex == 0 ? AnimationNames.BatOne : AnimationNames.BatTwo;
    }

    private string NextHammerAnimation()
    {
        swingIndex = (swingIndex + 1) % 2;
        return swingIndex == 0 ? AnimationNames.HammerOne : AnimationNames.HammerTwo;
    }

    private void ApplyVisibility()
    {
        // Toggle in-hand models based on current weapon
        if (batInHand != null)
        {
            batInHand.SetActive(CurrentWeapon == WeaponType.Bat);
            Debug.Log($"[WeaponManager] batInHand active={batInHand.activeSelf} (weapon={CurrentWeapon})");
        }
        else
        {
            Debug.LogWarning("[WeaponManager] 'batInHand' is not assigned in inspector.");
        }

        if (hammerInHand != null)
        {
            hammerInHand.SetActive(CurrentWeapon == WeaponType.Hammer);
            Debug.Log($"[WeaponManager] hammerInHand active={hammerInHand.activeSelf} (weapon={CurrentWeapon})");
        }
        else
        {
            Debug.LogWarning("[WeaponManager] 'hammerInHand' is not assigned in inspector.");
        }
    }

    // Keep the last pickup object used so we can re-enable it on drop (optional)
    private WeaponPickup lastPickupSource;

    public void EquipWeaponFromPickup(WeaponType weapon, WeaponPickup pickup)
    {
        lastPickupSource = pickup;
        EquipWeapon(weapon);
    }

    public void DropWeapon()
    {
        if (!IsWeaponEquipped) return;

        // Play the pickup animation even when dropping
        var anim = GetComponent<AnimationManager>();
        if (anim != null)
        {
            anim.PlayAnimation(AnimationNames.WeaponPickup);
        }

        if (lastPickupSource != null)
        {
            // Let the pickup place itself and avoid instant trigger
            lastPickupSource.PlaceAsDropped(transform);
            lastPickupSource = null;
        }

        // Block immediate re-equip for a short time
        _pickupCooldownUntil = Time.time + dropReequipBlockSeconds;

        UnequipWeapon();
    }

    private void OnPickupOrDrop(InputAction.CallbackContext ctx)
    {
        // Ignore if within cooldown window
        if (Time.time < _pickupCooldownUntil)
            return;

        // Requirement 3: If the character has a weapon in hand, pressing the same button drops it (and plays pickup anim)
        if (IsWeaponEquipped)
        {
            DropWeapon();
            return;
        }

        // Requirement 1 & 2: Only when in range and no weapon equipped, allow pickup and play pickup anim
        if (WeaponPickup.currentPickupInRange != null)
        {
            WeaponPickup.currentPickupInRange.TryEquip();
        }
        // Otherwise do nothing (no animation out of range)
    }
}