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

/// <summary>
/// CashSystemAI - AI decision layer for Cash System mode.
/// Works alongside CharacterAIController for combat.
/// </summary>
public class CashSystemAI : MonoBehaviour
{
    public enum AIState
    {
        Idle,
        CollectCash,
        DepositCash,
        AttackEnemy,
        StealFromVault,
        DefendVault,
        Retreat,
        Patrol        // NEW: wander when nothing to do
    }
    
    public enum AIRole
    {
        Balanced,
        Collector,
        Defender,
        Aggressor
    }
    
    [Header("AI Settings")]
    [SerializeField] private AIState currentState = AIState.Idle;
    [SerializeField] private AIRole role = AIRole.Balanced;
    
    [Header("Decision Weights")]
    [SerializeField] private float collectWeight = 1f;
    [SerializeField] private float depositWeight = 1.5f;
    [SerializeField] private float attackWeight = 1f;
    [SerializeField] private float stealWeight = 0.8f;
    [SerializeField] private float defendWeight = 1.2f;
    
    [Header("Thresholds")]
    [SerializeField] private float depositCashThreshold = 100f;
    [SerializeField] private float retreatHealthThreshold = 30f;
    [SerializeField] private float nearbyEnemyDistance = 15f;
    [SerializeField] private float nearbyBoxDistance = 50f;
    
    [Header("Patrol Settings")]
    [SerializeField] private float patrolRadius = 12f;
    [SerializeField] private float patrolPauseMin = 1f;
    [SerializeField] private float patrolPauseMax = 3f;
    
    [Header("Timing")]
    [SerializeField] private float decisionInterval = 0.3f;
    [SerializeField] private float stateMinDuration = 0.8f;
    [Header("Difficulty")]
    [SerializeField] private int aiLevel = 1; // 1 = normal, higher = harder/faster
    
    // Components
    private CashCarrier cashCarrier;
    private HealthNew health;
    private CharacterAIController aiController;
    private PlayerScript playerScript;
    
    // Runtime state
    private float lastDecisionTime;
    private float stateStartTime;
    private Transform currentTarget;
    private TeamVault myVault;
    private TeamVault enemyVault;
    private string myTeamTag;
    
    // Patrol state
    private Vector3 patrolDestination;
    private bool hasPatrolDestination;
    private float patrolPauseEndTime;
    private NavMeshAgent navAgent;
    
    public AIState CurrentState => currentState;
    public AIRole CurrentRole => role;
    
    private void Start()
    {
        // Only activate in Cash System mode
        if (SelectionOptions.Instance == null || !SelectionOptions.Instance.isCashSystemSelected)
        {
            enabled = false;
            return;
        }
        
        // Get components
        cashCarrier = GetComponent<CashCarrier>();
        health = GetComponent<HealthNew>();
        aiController = GetComponent<CharacterAIController>();
        playerScript = GetComponent<PlayerScript>();
        navAgent = GetComponent<NavMeshAgent>();
        
        myTeamTag = gameObject.tag;
        
        // Assign random role
        AssignRole();

        // Apply simple difficulty scaling
        if (aiLevel > 1)
        {
            float scale = 1f + (aiLevel - 1) * 0.15f;
            collectWeight *= scale;
            depositWeight *= scale;
            attackWeight *= scale;
            stealWeight *= scale;
            defendWeight *= scale;
            decisionInterval = Mathf.Max(0.2f, decisionInterval / (1f + (aiLevel - 1) * 0.2f));

            if (navAgent != null)
            {
                navAgent.speed *= 1f + (aiLevel - 1) * 0.1f;
            }
        }
        
        // Find vaults
        FindVaults();
        
        Debug.Log($"[CashSystemAI] Initialized on {gameObject.name}, Team: {myTeamTag}, Role: {role}");
    }
    
    private void AssignRole()
    {
        // Random role assignment with some bias based on character
        float rand = Random.value;
        if (rand < 0.4f) role = AIRole.Balanced;
        else if (rand < 0.65f) role = AIRole.Collector;
        else if (rand < 0.85f) role = AIRole.Defender;
        else role = AIRole.Aggressor;
        
        // Adjust weights based on role
        switch (role)
        {
            case AIRole.Collector:
                collectWeight *= 1.5f;
                depositWeight *= 1.3f;
                attackWeight *= 0.7f;
                break;
            case AIRole.Defender:
                defendWeight *= 2f;
                attackWeight *= 1.2f;
                collectWeight *= 0.6f;
                break;
            case AIRole.Aggressor:
                attackWeight *= 2f;
                stealWeight *= 1.5f;
                collectWeight *= 0.5f;
                break;
        }
    }
    
    private void FindVaults()
    {
        // Use EntityRegistry (zero-alloc) instead of FindObjectsOfType
        var vaults = EntityRegistry<TeamVault>.GetAll();
        for (int i = 0; i < vaults.Count; i++)
        {
            if (vaults[i] == null) continue;
            if (vaults[i].TeamTag == myTeamTag)
                myVault = vaults[i];
            else
                enemyVault = vaults[i];
        }
        
        // Fallback if EntityRegistry empty (vaults not yet registered)
        if (myVault == null || enemyVault == null)
        {
            TeamVault[] fallback = FindObjectsOfType<TeamVault>();
            foreach (TeamVault vault in fallback)
            {
                if (vault.TeamTag == myTeamTag) myVault = vault;
                else enemyVault = vault;
            }
        }
    }
    
    private void Update()
    {
        if (Time.time - lastDecisionTime < decisionInterval)
            return;
        
        lastDecisionTime = Time.time;
        
        // Lazy re-find vaults if they weren't found at Start
        if (myVault == null || enemyVault == null)
            FindVaults();
        
        // Check health for retreat (only when combat enabled)
        if (IsCombatEnabled() && health != null && health.currentHealth < retreatHealthThreshold)
        {
            SetState(AIState.Retreat);
            ExecuteState();
            return;
        }
        
        // Don't switch state too quickly
        if (Time.time - stateStartTime < stateMinDuration && currentState != AIState.Idle)
        {
            // But always re-execute current state so movement continues
            ExecuteState();
            return;
        }
        
        MakeDecision();
        ExecuteState();
    }
    
    /// <summary>
    /// Check if this AI is allowed to take combat actions (after countdown)
    /// </summary>
    private bool IsCombatEnabled()
    {
        // Check CharacterAIController's attack flag
        if (aiController != null)
        {
            return aiController.attack;
        }
        
        // Check PlayerScript's attack flag (for AI teammates)
        if (playerScript != null)
        {
            return playerScript.attack;
        }
        
        // Default to true if no scripts found
        return true;
    }
    
    private void MakeDecision()
    {
        float bestScore = 0f;
        AIState bestState = AIState.Idle;
        
        bool isCarrying = cashCarrier != null && cashCarrier.CarriedCash >= 0.01f;
        
        // Evaluate: Deposit (if carrying ANY cash and vault exists)
        if (isCarrying && myVault != null)
        {
            // Higher urgency the more cash we carry
            float score = depositWeight * Mathf.Max(1f, cashCarrier.CarriedCash / depositCashThreshold);
            // Extra urgency if carrying max bags
            if (cashCarrier.BagCount >= 3) score *= 2f;
            if (score > bestScore)
            {
                bestScore = score;
                bestState = AIState.DepositCash;
            }
        }
        
        // Evaluate: Defend (if vault under attack)
        if (myVault != null && myVault.IsBeingStolen)
        {
            float score = defendWeight * 3f; // High priority
            if (score > bestScore)
            {
                bestScore = score;
                bestState = AIState.DefendVault;
            }
        }
        
        // Evaluate: Attack (if enemy nearby AND combat enabled)
        if (IsCombatEnabled())
        {
            GameObject nearestEnemy = FindNearestEnemy();
            if (nearestEnemy != null)
            {
                float dist = Vector3.Distance(transform.position, nearestEnemy.transform.position);
                if (dist < nearbyEnemyDistance)
                {
                    float score = attackWeight * (1f - dist / nearbyEnemyDistance);
                    
                    // Higher score if enemy is carrying cash
                    CashCarrier enemyCarrier = nearestEnemy.GetComponent<CashCarrier>();
                    if (enemyCarrier != null && enemyCarrier.CarriedCash > 0)
                    {
                        score *= 1.5f;
                    }
                    
                    if (score > bestScore)
                    {
                        bestScore = score;
                        bestState = AIState.AttackEnemy;
                        currentTarget = nearestEnemy.transform;
                    }
                }
            }
        }
        
        // Evaluate: Collect (if not carrying max bags and bag available)
        if (cashCarrier != null && cashCarrier.CanPickUp())
        {
            CashBag nearestBag = FindNearestCashBag();
            if (nearestBag != null)
            {
                float dist = Vector3.Distance(transform.position, nearestBag.transform.position);
                // Always consider collecting if we can carry more
                float score = collectWeight * (1f + nearestBag.CashValue / 100f);
                // Slight distance penalty (prefer closer bags) but don't disqualify far bags
                score *= Mathf.Clamp01(1f - dist / (nearbyBoxDistance * 2f)) + 0.5f;
                
                if (score > bestScore)
                {
                    bestScore = score;
                    bestState = AIState.CollectCash;
                    currentTarget = nearestBag.transform;
                }
            }
        }
        
        // Evaluate: Steal (if enemy vault has cash and we're close and combat enabled)
        if (IsCombatEnabled() && enemyVault != null && enemyVault.StoredCash > 0 && !isCarrying)
        {
            float dist = Vector3.Distance(transform.position, enemyVault.transform.position);
            if (dist < nearbyBoxDistance)
            {
                float score = stealWeight * (enemyVault.StoredCash / 500f);
                if (score > bestScore)
                {
                    bestScore = score;
                    bestState = AIState.StealFromVault;
                    currentTarget = enemyVault.transform;
                }
            }
        }
        
        // Default: Patrol if nothing else to do (NEVER stay idle)
        if (bestState == AIState.Idle)
        {
            CashBag fallbackBag = FindNearestCashBag();
            if (fallbackBag != null)
            {
                bestState = AIState.CollectCash;
                currentTarget = fallbackBag.transform;
            }
            else
            {
                bestState = AIState.Patrol;
            }
        }
        
        SetState(bestState);
    }
    
    private void SetState(AIState newState)
    {
        if (newState != currentState)
        {
            currentState = newState;
            stateStartTime = Time.time;
            Debug.Log($"[CashSystemAI] {gameObject.name} -> {newState}");
        }
    }
    
    private void ExecuteState()
    {
        switch (currentState)
        {
            case AIState.Idle:
                ExecuteIdle();
                break;
            case AIState.CollectCash:
                ExecuteCollect();
                break;
            case AIState.DepositCash:
                ExecuteDeposit();
                break;
            case AIState.AttackEnemy:
                ExecuteAttack();
                break;
            case AIState.StealFromVault:
                ExecuteSteal();
                break;
            case AIState.DefendVault:
                ExecuteDefend();
                break;
            case AIState.Retreat:
                ExecuteRetreat();
                break;
            case AIState.Patrol:
                ExecutePatrol();
                break;
        }
    }
    
    private void ExecuteIdle()
    {
        // Try to find something to do
        CashBag nearestBag = FindNearestCashBag();
        if (nearestBag != null)
        {
            currentTarget = nearestBag.transform;
            SetState(AIState.CollectCash);
        }
        else
        {
            // Nothing to do — patrol instead of standing still
            SetState(AIState.Patrol);
        }
    }
    
    /// <summary>
    /// GTA-style patrol: pick random NavMesh point, walk there, pause briefly, repeat.
    /// This ensures AI NEVER stands completely still.
    /// </summary>
    private void ExecutePatrol()
    {
        // If pausing at a patrol point, wait
        if (Time.time < patrolPauseEndTime)
            return;
        
        // Check if we should switch to a real task
        CashBag nearestBag = FindNearestCashBag();
        if (nearestBag != null)
        {
            currentTarget = nearestBag.transform;
            SetState(AIState.CollectCash);
            hasPatrolDestination = false;
            return;
        }
        
        // Check for nearby enemies while patrolling
        GameObject nearestEnemy = FindNearestEnemy();
        if (nearestEnemy != null)
        {
            float dist = Vector3.Distance(transform.position, nearestEnemy.transform.position);
            if (dist < nearbyEnemyDistance * 0.5f) // React at half range while patrolling
            {
                currentTarget = nearestEnemy.transform;
                SetState(AIState.AttackEnemy);
                hasPatrolDestination = false;
                return;
            }
        }
        
        // Pick a new patrol destination if needed
        if (!hasPatrolDestination || HasReachedDestination())
        {
            if (hasPatrolDestination)
            {
                // Just arrived — pause briefly before picking a new point
                patrolPauseEndTime = Time.time + Random.Range(patrolPauseMin, patrolPauseMax);
                hasPatrolDestination = false;
                return;
            }
            
            // Pick random NavMesh point
            Vector3 randomPoint = transform.position + Random.insideUnitSphere * patrolRadius;
            randomPoint.y = transform.position.y;
            
            NavMeshHit hit;
            if (NavMesh.SamplePosition(randomPoint, out hit, patrolRadius, NavMesh.AllAreas))
            {
                patrolDestination = hit.position;
                hasPatrolDestination = true;
                MoveToward(patrolDestination);
            }
        }
        else
        {
            // Continue moving toward current patrol point
            MoveToward(patrolDestination);
        }
    }
    
    private bool HasReachedDestination()
    {
        if (!hasPatrolDestination) return true;
        float dist = Vector3.Distance(transform.position, patrolDestination);
        return dist < 2f;
    }
    
    private void ExecuteCollect()
    {
        if (currentTarget == null)
        {
            // Target gone — find another bag
            CashBag newBag = FindNearestCashBag();
            if (newBag != null)
            {
                currentTarget = newBag.transform;
            }
            else
            {
                SetState(AIState.Patrol);
                return;
            }
        }
        
        // Validate target is still an available bag
        var bag = currentTarget.GetComponent<CashBag>();
        if (bag == null || !bag.IsAvailable)
        {
            // Bag was taken or despawned — find another
            CashBag newBag = FindNearestCashBag();
            if (newBag != null)
            {
                currentTarget = newBag.transform;
                bag = newBag;
            }
            else
            {
                currentTarget = null;
                SetState(AIState.Patrol);
                return;
            }
        }
        
        float dist = Vector3.Distance(transform.position, currentTarget.position);
        
        // If we're close enough to the bag, attempt direct pickup
        if (dist <= 2.5f)
        {
            if (cashCarrier != null && cashCarrier.CanPickUp())
            {
                bag.PickUp(cashCarrier);
                cashCarrier.OnBagPickedUp(bag);
                
                Debug.Log($"[CashSystemAI] {gameObject.name} picked up bag (${bag.CashValue}), now carrying ${cashCarrier.CarriedCash:F0}");

                // After pickup, immediately decide next action
                if (cashCarrier.ShouldDeposit() && myVault != null)
                    SetState(AIState.DepositCash);
                else
                {
                    // Try to find another bag
                    CashBag nextBag = FindNearestCashBag();
                    if (nextBag != null && cashCarrier.CanPickUp())
                    {
                        currentTarget = nextBag.transform;
                        SetState(AIState.CollectCash);
                    }
                    else
                        SetState(AIState.DepositCash);
                }
                return;
            }
        }

        MoveToward(currentTarget.position);
    }
    
    private void ExecuteDeposit()
    {
        if (myVault == null)
        {
            FindVaults();
            if (myVault == null)
            {
                SetState(AIState.Patrol);
                return;
            }
        }
        
        // Check if we still have cash to deposit
        if (cashCarrier == null || cashCarrier.CarriedCash < 0.01f)
        {
            Debug.Log($"[CashSystemAI] {gameObject.name} has no cash to deposit, switching to collect");
            SetState(AIState.Idle);
            return;
        }
        
        float dist = Vector3.Distance(transform.position, myVault.transform.position);
        
        // If close enough to vault, deposit directly
        if (dist <= 3f)
        {
            float deposited = cashCarrier.DepositToVault(myVault);
            Debug.Log($"[CashSystemAI] {gameObject.name} deposited ${deposited:F0} to vault!");
            
            // Go find more bags after depositing
            SetState(AIState.Idle);
            return;
        }
        
        // Move toward vault
        MoveToward(myVault.transform.position);
    }
    
    private void ExecuteAttack()
    {
        if (currentTarget == null)
        {
            SetState(AIState.Idle);
            return;
        }
        
        // Let CharacterAIController handle combat
        if (aiController != null)
        {
            aiController.SetTarget(currentTarget.gameObject);
            aiController.NavigateTo(currentTarget.position);
        }
        
        if (aiController != null)
        {
            aiController.attack = true;
        }
    }
    
    private void ExecuteSteal()
    {
        if (enemyVault == null)
        {
            SetState(AIState.Idle);
            return;
        }
        
        MoveToward(enemyVault.transform.position);
    }
    
    private void ExecuteDefend()
    {
        if (myVault == null)
        {
            SetState(AIState.Idle);
            return;
        }
        
        // Find the thief attacking our vault
        GameObject nearestEnemy = FindNearestEnemy();
        if (nearestEnemy != null)
        {
            currentTarget = nearestEnemy.transform;
            
            if (aiController != null)
            {
                aiController.SetTarget(nearestEnemy);
                aiController.NavigateTo(nearestEnemy.transform.position);
            }
            
            if (aiController != null)
            {
                aiController.attack = true;
            }
        }
        else
        {
            MoveToward(myVault.transform.position);
        }
    }
    
    private void ExecuteRetreat()
    {
        if (myVault == null)
        {
            SetState(AIState.Idle);
            return;
        }
        
        // Run back to vault
        MoveToward(myVault.transform.position);
        
        // Deposit if we have cash
        if (cashCarrier != null && cashCarrier.CarriedCash > 0)
        {
            float dist = Vector3.Distance(transform.position, myVault.transform.position);
            if (dist < 3f)
            {
                cashCarrier.DepositToVault(myVault);
            }
        }
        
        // Recover health (if we have regen)
        if (health != null && health.currentHealth > retreatHealthThreshold * 1.5f)
        {
            SetState(AIState.Idle);
        }
    }
    
    private void MoveToward(Vector3 position)
    {
        if (aiController != null)
        {
            aiController.NavigateTo(position);
        }
        else
        {
            // Direct movement if no AI controller
            Vector3 dir = (position - transform.position).normalized;
            transform.position += dir * 3f * Time.deltaTime;
            transform.LookAt(position);
        }
    }
    
    private GameObject FindNearestEnemy()
    {
        // Use TeamMember static registry (zero-allocation) instead of FindGameObjectsWithTag
        GameObject nearest = null;
        float nearestDist = float.MaxValue;
        
        foreach (TeamMember member in TeamMember.AllMembers)
        {
            if (member == null || member.gameObject == gameObject) continue;
            if (member.gameObject.CompareTag(myTeamTag)) continue; // Same team, skip
            
            // Check if alive
            HealthNew enemyHealth = member.GetComponent<HealthNew>();
            if (enemyHealth != null && enemyHealth.currentHealth <= 0) continue;
            
            float dist = Vector3.Distance(transform.position, member.transform.position);
            if (dist < nearestDist)
            {
                nearestDist = dist;
                nearest = member.gameObject;
            }
        }
        
        return nearest;
    }
    
    private CashBag FindNearestCashBag()
    {
        // Use EntityRegistry (zero-alloc) instead of FindObjectsOfType
        var bags = EntityRegistry<CashBag>.GetAll();
        
        CashBag nearest = null;
        float nearestDist = float.MaxValue;
        
        for (int i = 0; i < bags.Count; i++)
        {
            CashBag bag = bags[i];
            if (bag == null) continue;
            // Accept Available or Dropped bags (not Carried)
            if (!bag.IsAvailable && !bag.IsDropped) continue;
            
            float dist = Vector3.Distance(transform.position, bag.transform.position);
            if (dist < nearestDist)
            {
                nearestDist = dist;
                nearest = bag;
            }
        }
        
        return nearest;
    }
}