using System.Collections.Generic;
using System.Linq;
using UnityEngine;

public class GridAI : MonoBehaviour
{
    GridBoard board;
    Team aiTeam;
    float thinkDelay = 0.5f;
    float thinkTimer = 0f;
    bool hasMoved = false;

    public void Initialize(GridBoard board, Team team)
    {
        this.board = board;
        this.aiTeam = team;
    }

    void Update()
    {
        if (board == null || board.IsGameOver()) return;
        if (board.GetCurrentTeam() != aiTeam)
        {
            hasMoved = false;
            thinkTimer = 0f;
            return;
        }

        if (hasMoved) return;

        // Small delay to make AI moves visible
        thinkTimer += Time.deltaTime;
        if (thinkTimer < thinkDelay) return;

        MakeMove();
        hasMoved = true;
    }

    void MakeMove()
    {
        var units = board.GetUnits(aiTeam);
        var visibleCells = board.GetVisibleCells(aiTeam);

        // Find a unit that can move
        GridUnit bestUnit = null;
        Vector2Int bestMove = default;
        float bestScore = float.MinValue;

        foreach (var unit in units)
        {
            if (unit.IsTaggedOut) continue;

            // Check speed nerf
            if (board.IsDefending(unit) && unit.ConsecutiveDefenseMoves % 4 == 3)
            {
                continue; // This unit must skip
            }

            var validMoves = board.GetValidMoves(unit);
            if (validMoves.Count == 0) continue;

            foreach (var move in validMoves)
            {
                float score = EvaluateMove(unit, move, visibleCells);
                if (score > bestScore)
                {
                    bestScore = score;
                    bestUnit = unit;
                    bestMove = move;
                }
            }
        }

        if (bestUnit != null)
        {
            board.ExecuteAIMove(bestUnit, bestMove);
        }
        else
        {
            // No valid moves, skip turn
            board.AISkipTurn();
        }
    }

    float EvaluateMove(GridUnit unit, Vector2Int move, HashSet<Vector2Int> visibleCells)
    {
        float score = 0f;

        var enemyFlagPos = board.GetEnemyFlagPosition(aiTeam);
        var ourFlagCarrier = board.GetFlagCarrier(aiTeam == Team.Blue ? Team.Red : Team.Blue);
        var theirFlagCarrier = board.GetFlagCarrier(aiTeam);

        // Are we carrying the enemy flag?
        bool carryingFlag = theirFlagCarrier == unit;

        if (carryingFlag)
        {
            // Priority: Return to base with flag
            // Move toward our defense zone
            int targetY = aiTeam == Team.Blue ? 0 : ZoneBoundaries.BoardHeight - 1;
            float distToBase = Mathf.Abs(move.y - targetY);
            score += 1000f - distToBase * 10f;
        }
        else if (ourFlagCarrier != null && visibleCells.Contains(ourFlagCarrier.GridPosition))
        {
            // Our flag is being carried - chase the carrier!
            float distToCarrier = ChebyshevDistance(move, ourFlagCarrier.GridPosition);
            score += 500f - distToCarrier * 15f;
        }
        else
        {
            // Go for the enemy flag
            float distToFlag = ChebyshevDistance(move, enemyFlagPos);
            score += 100f - distToFlag * 5f;
        }

        // Small bonus for advancing toward enemy
        if (aiTeam == Team.Blue)
        {
            score += move.y * 0.5f; // Blue advances up
        }
        else
        {
            score += (ZoneBoundaries.BoardHeight - move.y) * 0.5f; // Red advances down
        }

        // Avoid staying in defense too long (speed penalty)
        if (board.GetZoneOwner(move.y) == (aiTeam == Team.Blue ? ZoneOwner.Blue : ZoneOwner.Red))
        {
            score -= 5f;
        }

        // Small randomness to prevent predictability
        score += Random.Range(0f, 2f);

        return score;
    }

    int ChebyshevDistance(Vector2Int a, Vector2Int b)
    {
        return Mathf.Max(Mathf.Abs(a.x - b.x), Mathf.Abs(a.y - b.y));
    }
}