using UnityEngine;
using UnityEngine.AI;
using System.Collections;

public class Unit : MonoBehaviour
{
    [Header("Unit Info")]
    public string unitName = "Worker";
    public UnitType unitType = UnitType.Worker;
    public int maxHealth = 100;
    public float moveSpeed = 3.5f;
    public float rotationSpeed = 360f;
    
    [Header("Work Settings")]
    public float workRange = 2f;
    public float workDuration = 3f;
    public int resourceCapacity = 10;
    
    [Header("Combat (Optional)")]
    public int attackDamage = 0;
    public float attackRange = 0f;
    public float attackCooldown = 1f;
    
    [Header("UI")]
    public Transform healthBarPosition;
    public GameObject selectionRing;
    
    [Header("Role Changing")]
    public GameObject[] workerEquipment;  // Tools, bags, etc.
    public GameObject[] guardEquipment;   // Weapons, armor, etc.
    public GameObject[] builderEquipment; // Hammers, blueprints, etc.
    public GameObject[] scoutEquipment;   // Bow, cape, etc.
    
    [Header("Tool Management")]
    public Transform rightHandContainer; // The "R_hand_container" transform
    public GameObject axeTool; // For Wood resources
    public GameObject pickaxeTool; // For Stone resources  
    public GameObject sickleScytheTool; // For Food resources
    public GameObject goldPickTool; // For Gold resources
    
    // Private variables
    private int currentHealth;
    private NavMeshAgent navAgent;
    private Animator animator;
    private UnitState currentState = UnitState.Idle;
    private Vector3 targetPosition;
    public GameObject targetObject;
    private bool isSelected = false;
    public int currentResources = 0;
    public ResourceType currentResourceType = ResourceType.Wood; // Track what type of resource we're carrying
    private bool hasResources = false; // Track if we have any resources
    private bool isWorking = false;
    private float lastAttackTime;
    private GameObject lastWorkTarget; // Remember what we were working on
    private GameObject lastResourceTarget; // Remember the resource we were harvesting
    
    // Role changing variables
    private UnitType originalType;
    private UnitType temporaryType = UnitType.Worker; // Default to no temporary role
    private bool hasTemporaryRole = false;
    private float temporaryRoleTimer = 0f;
    private float temporaryRoleDuration = 0f;
    
    // Original stats for restoration
    private int originalAttackDamage;
    private float originalAttackRange;
    private float originalAttackCooldown;
    private float originalMoveSpeed;
    private int originalResourceCapacity;
    private float originalWorkRange;
    private float originalWorkDuration;
    
    // Animation parameters
    private static readonly int IsWalking = Animator.StringToHash("isWalking");
    private static readonly int IsRunning = Animator.StringToHash("isRunning");
    private static readonly int IsMining = Animator.StringToHash("isMining");
    private static readonly int IsChopping = Animator.StringToHash("isChopping");
    private static readonly int IsDead = Animator.StringToHash("isDead");
    private static readonly int TakeDamageParam = Animator.StringToHash("TakeDamage");
    private static readonly int DieParam = Animator.StringToHash("Die");
    
    // Events
    public System.Action<Unit> OnUnitDeath;
    public System.Action<Unit, int> OnResourcesChanged;
    public System.Action<Unit, UnitState> OnStateChanged;
    public System.Action<Unit, UnitType, UnitType> OnRoleChanged; // (unit, fromType, toType)
    public System.Action<Unit> OnTemporaryRoleExpired;
    
    // Tool Management variables
    private GameObject currentEquippedTool;
    private ResourceType currentRequiredTool = ResourceType.Wood;
    private bool hasRequiredTool = false;
    
    void Start()
    {
        InitializeUnit();
    }
    
    void InitializeUnit()
    {
        currentHealth = maxHealth;
        
        // Store original type and stats for restoration
        originalType = unitType;
        StoreOriginalStats();
        
        // Get NavMesh Agent
        navAgent = GetComponent<NavMeshAgent>();
        if (navAgent == null)
        {
            navAgent = gameObject.AddComponent<NavMeshAgent>();
        }
        navAgent.speed = moveSpeed;
        navAgent.angularSpeed = rotationSpeed;
        navAgent.stoppingDistance = 0.5f;
        
        // Get Animator
        animator = GetComponent<Animator>();
        
        // Setup selection ring
        if (selectionRing != null)
        {
            selectionRing.SetActive(false);
        }
        
        // Setup initial equipment
        UpdateEquipmentVisuals();
        
        ChangeState(UnitState.Idle);
    }
    
    void Update()
    {
        UpdateMovementAnimation();
        HandleStateBehavior();
        HandleTemporaryRole();
    }
    
    void HandleTemporaryRole()
    {
        if (hasTemporaryRole)
        {
            temporaryRoleTimer += Time.deltaTime;
            if (temporaryRoleTimer >= temporaryRoleDuration)
            {
                RevertToOriginalRole();
            }
        }
    }
    
    void UpdateMovementAnimation()
    {
        if (animator != null && navAgent != null)
        {
            float speed = navAgent.velocity.magnitude;
            bool isMoving = speed > 0.1f;
            bool isRunning = speed > 3f; // Running if speed is high
            bool isWalking = isMoving && !isRunning;
            
            // Set walking animation
            if (HasAnimatorParameter(animator, IsWalking))
            {
                animator.SetBool(IsWalking, isWalking);
            }
            
            // Set running animation
            if (HasAnimatorParameter(animator, IsRunning))
            {
                animator.SetBool(IsRunning, isRunning);
            }
        }
    }
    
    bool HasAnimatorParameter(Animator animator, int parameterHash)
    {
        if (animator == null) return false;
        
        for (int i = 0; i < animator.parameterCount; i++)
        {
            if (animator.GetParameter(i).nameHash == parameterHash)
            {
                return true;
            }
        }
        return false;
    }
    
    void HandleStateBehavior()
    {
        switch (currentState)
        {
            case UnitState.Moving:
                HandleMoving();
                break;
            case UnitState.Working:
                HandleWorking();
                break;
            case UnitState.Attacking:
                HandleAttacking();
                break;
            case UnitState.Returning:
                HandleReturning();
                break;
        }
    }
    
    void HandleMoving()
    {
        if (navAgent != null && !navAgent.pathPending)
        {
            if (navAgent.remainingDistance < 0.5f)
            {
                if (targetObject != null)
                {
                    // Check if we can work on the target
                    if (CanWorkOnTarget(targetObject))
                    {
                        // Remember the resource we're working on
                        if (targetObject.GetComponent<ResourceNode>() != null)
                        {
                            lastResourceTarget = targetObject;
                            
                            // Make sure we have the right tool equipped
                            ResourceNode resourceNode = targetObject.GetComponent<ResourceNode>();
                            ResourceType resourceType = resourceNode.GetResourceType();
                            GameObject requiredTool = GetRequiredToolForResource(resourceType);
                            
                            if (requiredTool != null && currentEquippedTool != requiredTool)
                            {
                                EquipTool(requiredTool);
                                Debug.Log($"{unitName} equipped {requiredTool.name} for {resourceType}");
                            }
                        }
                        ChangeState(UnitState.Working);
                    }
                    else if (CanAttackTarget(targetObject))
                    {
                        ChangeState(UnitState.Attacking);
                    }
                    else
                    {
                        ChangeState(UnitState.Idle);
                    }
                }
                else
                {
                    ChangeState(UnitState.Idle);
                }
            }
        }
    }
    
    void HandleWorking()
    {
        if (targetObject == null)
        {
            ChangeState(UnitState.Idle);
            return;
        }
        
        if (!isWorking && IsInRange(targetObject, workRange))
        {
            StartCoroutine(WorkRoutine());
        }
        else if (!IsInRange(targetObject, workRange))
        {
            // Find a good position around the resource to avoid stacking
            Vector3 workPosition = FindWorkPositionAroundTarget(targetObject);
            MoveTo(workPosition, targetObject);
        }
    }
    
    void HandleAttacking()
    {
        if (targetObject == null)
        {
            ChangeState(UnitState.Idle);
            return;
        }
        
        if (IsInRange(targetObject, attackRange))
        {
            if (Time.time - lastAttackTime >= attackCooldown)
            {
                PerformAttack();
                lastAttackTime = Time.time;
            }
        }
        else
        {
            MoveTo(targetObject.transform.position, targetObject);
        }
    }
    
    void HandleReturning()
    {
        // Logic for returning to base/storage
        if (navAgent != null && !navAgent.pathPending && navAgent.remainingDistance < 0.5f)
        {
            // Check if we're at storage to deposit resources or get tools
            if (currentResources > 0)
            {
                Debug.Log($"{unitName} reached storage! Depositing {currentResources} {currentResourceType}...");
                DropOffResources();
            }
            
            // Check if we need to get a tool from storage
            if (NeedsToGetTool(currentRequiredTool))
            {
                Debug.Log($"{unitName} getting tool for {currentRequiredTool} from storage...");
                GameObject requiredTool = GetRequiredToolForResource(currentRequiredTool);
                if (requiredTool != null)
                {
                    EquipTool(requiredTool);
                    hasRequiredTool = true;
                    Debug.Log($"{unitName} equipped tool for {currentRequiredTool}");
                }
                else
                {
                    Debug.LogWarning($"{unitName} couldn't find required tool for {currentRequiredTool}!");
                }
            }
            
            // Check if we have a remembered resource to return to
            if (lastResourceTarget != null)
            {
                ResourceNode resourceNode = lastResourceTarget.GetComponent<ResourceNode>();
                if (resourceNode != null && resourceNode.GetResourceAmount() > 0)
                {
                    Debug.Log($"{unitName} returning to continue harvesting {resourceNode.GetResourceType()}...");
                    MoveTo(lastResourceTarget.transform.position, lastResourceTarget);
                    ChangeState(UnitState.Working);
                }
                else
                {
                    Debug.Log($"{unitName} previous resource depleted, going idle.");
                    lastResourceTarget = null;
                    ChangeState(UnitState.Idle);
                }
            }
            else
            {
                Debug.Log($"{unitName} no previous resource to return to, going idle.");
                ChangeState(UnitState.Idle);
            }
        }
    }
    
    IEnumerator WorkRoutine()
    {
        isWorking = true;
        
        // Face the target
        if (targetObject != null)
        {
            Vector3 direction = (targetObject.transform.position - transform.position).normalized;
            if (direction != Vector3.zero)
            {
                transform.rotation = Quaternion.LookRotation(direction);
            }
        }
        
        // Play work animation based on target type
        if (animator != null && targetObject != null)
        {
            if (targetObject.CompareTag("Rock") || targetObject.CompareTag("Mine"))
            {
                if (HasAnimatorParameter(animator, IsMining))
                {
                    animator.SetBool(IsMining, true);
                }
            }
            else if (targetObject.CompareTag("Tree") || targetObject.CompareTag("Wood"))
            {
                if (HasAnimatorParameter(animator, IsChopping))
                {
                    animator.SetBool(IsChopping, true);
                }
            }
        }
        
        yield return new WaitForSeconds(workDuration);
        
        // Stop work animation
        if (animator != null)
        {
            if (HasAnimatorParameter(animator, IsMining))
            {
                animator.SetBool(IsMining, false);
            }
            if (HasAnimatorParameter(animator, IsChopping))
            {
                animator.SetBool(IsChopping, false);
            }
        }
        
        // Collect resources or perform work
        if (targetObject != null)
        {
            PerformWork();
        }
        
        isWorking = false;
    }
    
    void PerformWork()
    {
        // Don't perform work if we're in returning state (at storage)
        if (currentState == UnitState.Returning)
        {
            return;
        }
        
        // Check if target has resources component
        ResourceNode resourceNode = targetObject.GetComponent<ResourceNode>();
        if (resourceNode != null && currentResources < resourceCapacity)
        {
            int gathered = resourceNode.GatherResource(1);
            if (gathered > 0)
            {
                // Set resource type if this is our first resource
                if (!hasResources)
                {
                    currentResourceType = resourceNode.GetResourceType();
                    hasResources = true;
                }
                
                currentResources += gathered;
                OnResourcesChanged?.Invoke(this, currentResources);
                
                // Debug log for player feedback
                Debug.Log($"{unitName} gathered {gathered} {resourceNode.GetResourceType()}! Total carried: {currentResources}/{resourceCapacity}");
                
                // Visual feedback - make unit briefly flash or scale
                StartCoroutine(GatheringFeedback());
            }
            
            // Check if inventory is full
            if (currentResources >= resourceCapacity)
            {
                Debug.Log($"{unitName} inventory full! Returning to storage...");
                // Find nearest storage and return
                GameObject storage = FindNearestStorage();
                if (storage != null)
                {
                    MoveTo(storage.transform.position, storage);
                    ChangeState(UnitState.Returning);
                }
                else
                {
                    Debug.LogWarning($"{unitName} couldn't find storage to return resources!");
                    ChangeState(UnitState.Idle);
                }
            }
            else if (resourceNode.GetResourceAmount() > 0)
            {
                // Continue working on the same resource - ensure animation plays again
                Debug.Log($"{unitName} continuing to harvest {resourceNode.GetResourceType()}...");
                
                // Small delay before starting next work cycle to ensure animation reset
                StartCoroutine(DelayedWorkContinue());
            }
            else
            {
                Debug.Log($"{unitName} finished gathering from {resourceNode.name} - resource depleted.");
                ChangeState(UnitState.Idle);
            }
        }
        else
        {
            if (resourceNode == null)
            {
                // Check if we're at storage - if so, just return without warning
                if (targetObject != null && (targetObject.CompareTag("Storage") || targetObject.CompareTag("StoreRoom")))
                {
                    Debug.Log($"{unitName} at storage, ready to deposit resources.");
                    return;
                }
                Debug.LogWarning($"{unitName} tried to work on object without ResourceNode component!");
            }
            else if (currentResources >= resourceCapacity)
            {
                Debug.Log($"{unitName} is already at full capacity!");
            }
            ChangeState(UnitState.Idle);
        }
    }
    
    IEnumerator GatheringFeedback()
    {
        // Scale up briefly for visual feedback
        Vector3 originalScale = transform.localScale;
        transform.localScale = originalScale * 1.1f;
        
        yield return new WaitForSeconds(0.1f);
        
        transform.localScale = originalScale;
    }
    
    void PerformAttack()
    {
        if (animator != null)
        {
            if (HasAnimatorParameter(animator, TakeDamageParam))
            {
                animator.SetTrigger(TakeDamageParam);
            }
        }
        
        // Deal damage to target
        Unit targetUnit = targetObject.GetComponent<Unit>();
        if (targetUnit != null)
        {
            targetUnit.TakeDamage(attackDamage);
        }
    }
    
    GameObject FindNearestStorage()
    {
        // Look for both "Storage" and "StoreRoom" tagged objects
        GameObject[] storages = GameObject.FindGameObjectsWithTag("Storage");
        GameObject[] storeRooms = GameObject.FindGameObjectsWithTag("StoreRoom");
        
        GameObject nearest = null;
        float nearestDistance = float.MaxValue;
        
        // Check Storage tagged objects
        foreach (GameObject storage in storages)
        {
            float distance = Vector3.Distance(transform.position, storage.transform.position);
            if (distance < nearestDistance)
            {
                nearestDistance = distance;
                nearest = storage;
            }
        }
        
        // Check StoreRoom tagged objects
        foreach (GameObject storeRoom in storeRooms)
        {
            float distance = Vector3.Distance(transform.position, storeRoom.transform.position);
            if (distance < nearestDistance)
            {
                nearestDistance = distance;
                nearest = storeRoom;
            }
        }
        
        if (nearest == null)
        {
            Debug.LogWarning($"{unitName} could not find any Storage or StoreRoom objects!");
        }
        
        return nearest;
    }
    
    void DropOffResources()
    {
        if (currentResources > 0)
        {
            // Find storage component and deposit resources
            Storage storage = targetObject?.GetComponent<Storage>();
            if (storage != null)
            {
                Debug.Log($"{unitName} depositing {currentResources} {currentResourceType} to storage.");
                bool success = storage.AddResources(currentResourceType, currentResources);
                if (success)
                {
                    currentResources = 0;
                    hasResources = false;
                    OnResourcesChanged?.Invoke(this, currentResources);
                    Debug.Log($"{unitName} successfully deposited resources.");
                }
                else
                {
                    Debug.LogWarning($"{unitName} couldn't deposit resources - storage full or doesn't accept {currentResourceType}!");
                }
            }
            else
            {
                Debug.LogWarning($"{unitName} couldn't find storage component on target!");
            }
        }
    }
    
    void ChangeState(UnitState newState)
    {
        currentState = newState;
        OnStateChanged?.Invoke(this, newState);
    }
    
    bool IsInRange(GameObject target, float range)
    {
        return Vector3.Distance(transform.position, target.transform.position) <= range;
    }
    
    // Public methods for external control
    public void MoveTo(Vector3 position, GameObject target = null)
    {
        if (navAgent != null)
        {
            targetPosition = position;
            targetObject = target;
            navAgent.SetDestination(position);
            ChangeState(UnitState.Moving);
        }
    }
    
    public void WorkOn(GameObject target)
    {
        if (CanWorkOnTarget(target))
        {
            ResourceNode targetResource = target.GetComponent<ResourceNode>();
            if (targetResource != null)
            {
                ResourceType targetResourceType = targetResource.GetResourceType();
                
                // Check if we need the right tool for this resource
                if (NeedsToGetTool(targetResourceType))
                {
                    Debug.Log($"{unitName} needs tool for {targetResourceType}. Going to storage first.");
                    
                    // Find storage to get the tool
                    GameObject storage = FindNearestStorage();
                    if (storage != null)
                    {
                        lastWorkTarget = target; // Remember the work target
                        lastResourceTarget = target; // Remember the resource target
                        currentRequiredTool = targetResourceType; // Remember what tool we need
                        MoveTo(storage.transform.position, storage);
                        ChangeState(UnitState.Returning); // Use returning state to get tool
                        return;
                    }
                    else
                    {
                        Debug.LogWarning($"{unitName} couldn't find storage to get tool!");
                        return;
                    }
                }
                
                // Check if we're carrying different resource type
                if (hasResources && currentResourceType != targetResourceType)
                {
                    Debug.Log($"{unitName} is carrying {currentResourceType} but target is {targetResourceType}. Returning to storage first.");
                    // Return to storage first, then work on new target
                    GameObject storage = FindNearestStorage();
                    if (storage != null)
                    {
                        lastWorkTarget = target; // Remember the new target
                        lastResourceTarget = target; // Remember the resource target
                        MoveTo(storage.transform.position, storage);
                        ChangeState(UnitState.Returning);
                        return;
                    }
                }
            }
            
            targetObject = target;
            lastWorkTarget = target;
            lastResourceTarget = target; // Also remember as the resource target
            
            // Find a good position around the target to avoid stacking
            Vector3 workPosition = FindWorkPositionAroundTarget(target);
            MoveTo(workPosition, target);
        }
    }
    
    public void AttackTarget(GameObject target)
    {
        if (CanAttackTarget(target))
        {
            targetObject = target;
            MoveTo(target.transform.position, target);
        }
    }
    
    public void SetSelected(bool selected)
    {
        isSelected = selected;
        if (selectionRing != null)
        {
            selectionRing.SetActive(selected);
        }
    }
    
    public void TakeDamage(int damage)
    {
        currentHealth -= damage;
        
        // Play damage animation
        if (animator != null)
        {
            if (HasAnimatorParameter(animator, TakeDamageParam))
            {
                animator.SetTrigger(TakeDamageParam);
            }
        }
        
        if (currentHealth <= 0)
        {
            Die();
        }
    }
    
    public void Heal(int amount)
    {
        currentHealth = Mathf.Min(currentHealth + amount, maxHealth);
    }
    
    void Die()
    {
        if (animator != null)
        {
            if (HasAnimatorParameter(animator, DieParam))
            {
                animator.SetTrigger(DieParam);
            }
            if (HasAnimatorParameter(animator, IsDead))
            {
                animator.SetBool(IsDead, true);
            }
        }
        
        ChangeState(UnitState.Dead);
        OnUnitDeath?.Invoke(this);
        
        // Disable components
        if (navAgent != null)
        {
            navAgent.enabled = false;
        }
        
        // Destroy after animation
        Destroy(gameObject, 2f);
    }
    
    // Role changing methods
    public void ChangeRolePermanently(UnitType newType)
    {
        if (newType == unitType) return;
        
        UnitType previousType = unitType;
        
        // If we had a temporary role, clear it first
        if (hasTemporaryRole)
        {
            hasTemporaryRole = false;
            temporaryRoleTimer = 0f;
        }
        
        // Change the permanent type
        originalType = newType;
        unitType = newType;
        unitName = GetUnitTypeName(newType);
        
        // Update stats and equipment
        ApplyRoleStats(newType);
        UpdateEquipmentVisuals();
        
        OnRoleChanged?.Invoke(this, previousType, newType);
        
        Debug.Log($"{gameObject.name} permanently changed from {previousType} to {newType}");
    }
    
    public void ChangeRoleTemporarily(UnitType newType, float duration)
    {
        if (newType == GetCurrentEffectiveType()) return;
        
        UnitType previousType = GetCurrentEffectiveType();
        
        // Set temporary role
        temporaryType = newType;
        hasTemporaryRole = true;
        temporaryRoleTimer = 0f;
        temporaryRoleDuration = duration;
        
        // Apply temporary stats and equipment
        ApplyRoleStats(newType);
        UpdateEquipmentVisuals();
        
        OnRoleChanged?.Invoke(this, previousType, newType);
        
        Debug.Log($"{gameObject.name} temporarily changed from {previousType} to {newType} for {duration} seconds");
    }
    
    public void RevertToOriginalRole()
    {
        if (!hasTemporaryRole) return;
        
        UnitType previousType = temporaryType;
        
        // Clear temporary role
        hasTemporaryRole = false;
        temporaryRoleTimer = 0f;
        temporaryType = UnitType.Worker; // Reset to default
        
        // Revert to original stats and equipment
        ApplyRoleStats(originalType);
        UpdateEquipmentVisuals();
        
        OnTemporaryRoleExpired?.Invoke(this);
        OnRoleChanged?.Invoke(this, previousType, originalType);
        
        Debug.Log($"{gameObject.name} reverted from {previousType} to {originalType}");
    }
    
    public UnitType GetCurrentEffectiveType()
    {
        return hasTemporaryRole ? temporaryType : unitType;
    }
    
    public bool HasTemporaryRole()
    {
        return hasTemporaryRole;
    }
    
    public float GetTemporaryRoleTimeRemaining()
    {
        return hasTemporaryRole ? (temporaryRoleDuration - temporaryRoleTimer) : 0f;
    }
    
    void StoreOriginalStats()
    {
        originalAttackDamage = attackDamage;
        originalAttackRange = attackRange;
        originalAttackCooldown = attackCooldown;
        originalMoveSpeed = moveSpeed;
        originalResourceCapacity = resourceCapacity;
        originalWorkRange = workRange;
        originalWorkDuration = workDuration;
    }
    
    void ApplyRoleStats(UnitType roleType)
    {
        switch (roleType)
        {
            case UnitType.Worker:
                attackDamage = 0;
                attackRange = 0f;
                attackCooldown = 1f;
                moveSpeed = 3.5f;
                resourceCapacity = 10;
                workRange = 2f;
                workDuration = 3f;
                break;
                
            case UnitType.Guard:
                attackDamage = 25;
                attackRange = 1.5f;
                attackCooldown = 2f;
                moveSpeed = 4f;
                resourceCapacity = 0;
                workRange = 0f;
                workDuration = 0f;
                break;
                
            case UnitType.Builder:
                attackDamage = 5;
                attackRange = 1f;
                attackCooldown = 3f;
                moveSpeed = 3f;
                resourceCapacity = 5;
                workRange = 3f;
                workDuration = 5f;
                break;
                
            case UnitType.Scout:
                attackDamage = 15;
                attackRange = 8f;
                attackCooldown = 1.5f;
                moveSpeed = 5f;
                resourceCapacity = 0;
                workRange = 0f;
                workDuration = 0f;
                break;
        }
        
        // Update NavMeshAgent speed
        if (navAgent != null)
        {
            navAgent.speed = moveSpeed;
        }
    }
    
    void UpdateEquipmentVisuals()
    {
        UnitType currentType = GetCurrentEffectiveType();
        
        // Hide all equipment first
        SetEquipmentActive(workerEquipment, false);
        SetEquipmentActive(guardEquipment, false);
        SetEquipmentActive(builderEquipment, false);
        SetEquipmentActive(scoutEquipment, false);
        
        // Show appropriate equipment
        switch (currentType)
        {
            case UnitType.Worker:
                SetEquipmentActive(workerEquipment, true);
                break;
            case UnitType.Guard:
                SetEquipmentActive(guardEquipment, true);
                break;
            case UnitType.Builder:
                SetEquipmentActive(builderEquipment, true);
                break;
            case UnitType.Scout:
                SetEquipmentActive(scoutEquipment, true);
                break;
        }
    }
    
    void SetEquipmentActive(GameObject[] equipment, bool active)
    {
        if (equipment == null) return;
        
        foreach (GameObject item in equipment)
        {
            if (item != null)
            {
                item.SetActive(active);
            }
        }
    }
    
    string GetUnitTypeName(UnitType type)
    {
        return type switch
        {
            UnitType.Worker => "Worker",
            UnitType.Guard => "Guard",
            UnitType.Builder => "Builder",
            UnitType.Scout => "Scout",
            _ => "Unknown"
        };
    }
    
    // Updated methods to use current effective type
    public bool CanWorkOnTarget(GameObject target)
    {
        UnitType currentType = GetCurrentEffectiveType();
        return target.GetComponent<ResourceNode>() != null && 
               (currentType == UnitType.Worker || currentType == UnitType.Builder);
    }
    
    bool CanAttackTarget(GameObject target)
    {
        return target.GetComponent<Unit>() != null && attackDamage > 0;
    }
    
    // Tool Management Methods
    GameObject GetRequiredToolForResource(ResourceType resourceType)
    {
        return resourceType switch
        {
            ResourceType.Wood => axeTool,
            ResourceType.Stone => pickaxeTool,
            ResourceType.Food => sickleScytheTool,
            ResourceType.Gold => goldPickTool,
            _ => null
        };
    }
    
    bool HasRequiredToolForResource(ResourceType resourceType)
    {
        GameObject requiredTool = GetRequiredToolForResource(resourceType);
        return requiredTool != null;
    }
    
    void EquipTool(GameObject tool)
    {
        if (rightHandContainer == null)
        {
            Debug.LogWarning($"{unitName} has no right hand container assigned!");
            return;
        }
        
        // Deactivate all tools first
        DeactivateAllTools();
        
        // Activate the required tool
        if (tool != null)
        {
            tool.SetActive(true);
            currentEquippedTool = tool;
            Debug.Log($"{unitName} equipped {tool.name}");
        }
    }
    
    void DeactivateAllTools()
    {
        if (rightHandContainer == null) return;
        
        // Deactivate all child objects in the right hand container
        for (int i = 0; i < rightHandContainer.childCount; i++)
        {
            rightHandContainer.GetChild(i).gameObject.SetActive(false);
        }
        
        currentEquippedTool = null;
    }
    
    bool NeedsToGetTool(ResourceType resourceType)
    {
        GameObject requiredTool = GetRequiredToolForResource(resourceType);
        
        // Check if we don't have the tool, or if we have the wrong tool equipped
        return requiredTool == null || currentEquippedTool != requiredTool;
    }
    
    // Getters
    public UnitState GetState() => currentState;
    public int GetHealth() => currentHealth;
    public int GetMaxHealth() => maxHealth;
    public int GetResources() => currentResources;
    public ResourceType GetCurrentResourceType() => currentResourceType;
    public bool HasResources() => hasResources;
    public bool IsSelected() => isSelected;
    public UnitType GetUnitType() => GetCurrentEffectiveType(); // Updated to return effective type
    public UnitType GetOriginalType() => originalType;
    
    Vector3 FindWorkPositionAroundTarget(GameObject target)
    {
        if (target == null) return transform.position;
        
        Vector3 targetPos = target.transform.position;
        float radius = workRange * 0.8f; // Work slightly inside the work range
        
        // Try to find an unoccupied position around the target
        for (int attempts = 0; attempts < 8; attempts++)
        {
            float angle = (360f / 8f) * attempts * Mathf.Deg2Rad;
            Vector3 testPosition = targetPos + new Vector3(
                Mathf.Cos(angle) * radius,
                0f,
                Mathf.Sin(angle) * radius
            );
            
            // Check if this position is free of other units
            Collider[] nearbyUnits = Physics.OverlapSphere(testPosition, 1f);
            bool positionFree = true;
            
            foreach (Collider col in nearbyUnits)
            {
                Unit otherUnit = col.GetComponent<Unit>();
                if (otherUnit != null && otherUnit != this)
                {
                    positionFree = false;
                    break;
                }
            }
            
            if (positionFree)
            {
                return testPosition;
            }
        }
        
        // If no free position found, use a random position around the target
        float randomAngle = Random.Range(0f, 360f) * Mathf.Deg2Rad;
        return targetPos + new Vector3(
            Mathf.Cos(randomAngle) * radius,
            0f,
            Mathf.Sin(randomAngle) * radius
        );
    }
    
    IEnumerator DelayedWorkContinue()
    {
        // Small delay to ensure animation state is properly reset
        yield return new WaitForSeconds(0.1f);
        
        // Start the next work cycle
        if (targetObject != null && currentState == UnitState.Working)
        {
            StartCoroutine(WorkRoutine());
        }
    }
}
public enum UnitType
{
    Worker,
    Guard,
    Builder,
    Scout
}

public enum UnitState
{
    Idle,
    Moving,
    Working,
    Attacking,
    Returning,
    Dead
}