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

public enum ZoneOwner { Blue, Neutral, Red }

public static class ZoneBoundaries
{
    public const int TeamBlueDefenseEnd = 20;  // Y < 20 is Blue defense
    public const int NeutralEnd = 30;           // Y < 30 is neutral (if >= TeamBlueDefenseEnd)
    public const int BoardWidth = 20;
    public const int BoardHeight = 50;
}

public class GridBoard : MonoBehaviour
{
    // Configuration
    const float CellSize = 1f;
    const float CellPadding = 0.05f;
    const int UnitsPerTeam = 3;
    const int WinScore = 3;
    const int RespawnDelay = 2;  // Turns until respawn
    const int VisionRadius = 3;  // Chebyshev distance for fog of war

    // Colors
    static readonly Color BlueZoneColor = new Color(0.2f, 0.3f, 0.6f, 0.4f);
    static readonly Color NeutralZoneColor = new Color(0.4f, 0.4f, 0.4f, 0.4f);
    static readonly Color RedZoneColor = new Color(0.6f, 0.2f, 0.2f, 0.4f);
    static readonly Color BlueUnitColor = new Color(0.3f, 0.5f, 1f);
    static readonly Color RedUnitColor = new Color(1f, 0.3f, 0.3f);
    static readonly Color ValidMoveColor = new Color(0.3f, 1f, 0.3f, 0.5f);
    static readonly Color SelectedColor = new Color(1f, 1f, 0.3f, 0.8f);
    static readonly Color FlagBlueColor = new Color(0.3f, 0.5f, 1f);
    static readonly Color FlagRedColor = new Color(1f, 0.3f, 0.3f);

    // Board state
    GridUnit[,] cellOccupants;
    Dictionary<GridUnit, Vector2Int> unitPositions = new();
    List<GridUnit> blueUnits = new();
    List<GridUnit> redUnits = new();

    // Flags
    Vector2Int blueFlagPosition;
    Vector2Int redFlagPosition;
    Vector2Int blueFlagHome;
    Vector2Int redFlagHome;
    GridUnit blueFlagCarrier;
    GridUnit redFlagCarrier;
    GameObject blueFlagGO;
    GameObject redFlagGO;

    // Turn state
    Team currentTeam = Team.Blue;
    int turnNumber = 0;

    // Scores
    int blueScore = 0;
    int redScore = 0;
    bool gameOver = false;

    // UI state
    GridUnit selectedUnit;
    List<Vector2Int> validMoves = new();
    List<GameObject> validMoveHighlights = new();
    GameObject selectionHighlight;

    // Visual objects
    GameObject[,] cellVisuals;

    // Visibility
    HashSet<Vector2Int> visibleToBlue = new();
    HashSet<Vector2Int> visibleToRed = new();

    // AI
    GridAI ai;

    // UI references (set by Game.cs)
    public TMPro.TextMeshProUGUI scoreText;
    public TMPro.TextMeshProUGUI gameOverText;
    public TMPro.TextMeshProUGUI turnText;

    int nextUnitId = 0;

    void Start()
    {
        InitializeBoard();
        CreateBoardVisuals();
        SpawnUnits();
        CreateFlags();
        RecalculateVisibility();
        UpdateUI();

        // Initialize AI for red team
        ai = gameObject.AddComponent<GridAI>();
        ai.Initialize(this, Team.Red);
    }

    void InitializeBoard()
    {
        cellOccupants = new GridUnit[ZoneBoundaries.BoardWidth, ZoneBoundaries.BoardHeight];
        cellVisuals = new GameObject[ZoneBoundaries.BoardWidth, ZoneBoundaries.BoardHeight];
    }

    void CreateBoardVisuals()
    {
        // Create parent for organization
        var boardParent = new GameObject("Board");
        boardParent.transform.SetParent(transform);

        for (int x = 0; x < ZoneBoundaries.BoardWidth; x++)
        {
            for (int y = 0; y < ZoneBoundaries.BoardHeight; y++)
            {
                Color zoneColor = GetZoneColor(y);
                var cell = CreateSprite($"Cell_{x}_{y}", zoneColor, CellSize - CellPadding, CellSize - CellPadding);
                cell.transform.SetParent(boardParent.transform);
                cell.transform.position = GridToWorld(new Vector2Int(x, y));
                cell.GetComponent<SpriteRenderer>().sortingOrder = -10;
                cellVisuals[x, y] = cell;
            }
        }
    }

    Color GetZoneColor(int y)
    {
        var zone = GetZoneOwner(y);
        return zone switch
        {
            ZoneOwner.Blue => BlueZoneColor,
            ZoneOwner.Neutral => NeutralZoneColor,
            ZoneOwner.Red => RedZoneColor,
            _ => NeutralZoneColor
        };
    }

    void SpawnUnits()
    {
        // Blue units spawn at bottom of blue zone
        for (int i = 0; i < UnitsPerTeam; i++)
        {
            int x = ZoneBoundaries.BoardWidth / 2 - 1 + (i % 3);
            int y = 2 + (i / 3);
            SpawnUnit(Team.Blue, new Vector2Int(x, y));
        }

        // Red units spawn at top of red zone
        for (int i = 0; i < UnitsPerTeam; i++)
        {
            int x = ZoneBoundaries.BoardWidth / 2 - 1 + (i % 3);
            int y = ZoneBoundaries.BoardHeight - 3 - (i / 3);
            SpawnUnit(Team.Red, new Vector2Int(x, y));
        }
    }

    GridUnit SpawnUnit(Team team, Vector2Int position)
    {
        Color color = team == Team.Blue ? BlueUnitColor : RedUnitColor;
        var go = CreateSprite($"{team}Unit_{nextUnitId}", color, CellSize * 0.8f, CellSize * 0.8f);
        go.GetComponent<SpriteRenderer>().sortingOrder = 10;

        var unit = new GridUnit(nextUnitId++, team, position, go);
        unit.SetWorldPosition(GridToWorld(position));

        cellOccupants[position.x, position.y] = unit;
        unitPositions[unit] = position;

        if (team == Team.Blue)
            blueUnits.Add(unit);
        else
            redUnits.Add(unit);

        return unit;
    }

    void CreateFlags()
    {
        // Blue flag at center bottom of blue zone
        blueFlagHome = new Vector2Int(ZoneBoundaries.BoardWidth / 2, 1);
        blueFlagPosition = blueFlagHome;
        blueFlagGO = CreateSprite("BlueFlag", FlagBlueColor, CellSize * 0.5f, CellSize * 0.8f);
        blueFlagGO.GetComponent<SpriteRenderer>().sortingOrder = 5;
        blueFlagGO.transform.position = GridToWorld(blueFlagPosition);

        // Red flag at center top of red zone
        redFlagHome = new Vector2Int(ZoneBoundaries.BoardWidth / 2, ZoneBoundaries.BoardHeight - 2);
        redFlagPosition = redFlagHome;
        redFlagGO = CreateSprite("RedFlag", FlagRedColor, CellSize * 0.5f, CellSize * 0.8f);
        redFlagGO.GetComponent<SpriteRenderer>().sortingOrder = 5;
        redFlagGO.transform.position = GridToWorld(redFlagPosition);
    }

    void Update()
    {
        if (gameOver) return;

        // AI handles red team
        if (currentTeam == Team.Red)
        {
            return; // AI takes control in GridAI.Update()
        }

        HandleInput();
    }

    void HandleInput()
    {
        var mouse = Mouse.current;
        if (mouse == null) return;

        if (mouse.leftButton.wasPressedThisFrame)
        {
            Vector2 worldPos = Camera.main.ScreenToWorldPoint(mouse.position.ReadValue());
            Vector2Int gridPos = WorldToGrid(worldPos);

            if (!IsInBounds(gridPos)) return;

            // If we have a unit selected and clicked on a valid move, execute it
            if (selectedUnit != null && validMoves.Contains(gridPos))
            {
                ExecutePlayerMove(selectedUnit, gridPos);
                return;
            }

            // Check if we clicked on a friendly unit to select it
            var unitAtCell = cellOccupants[gridPos.x, gridPos.y];
            if (unitAtCell != null && unitAtCell.Team == currentTeam && !unitAtCell.IsTaggedOut)
            {
                // Check if unit can move this turn (defense speed nerf)
                if (IsDefending(unitAtCell) && unitAtCell.ConsecutiveDefenseMoves % 4 == 3)
                {
                    // This unit must skip their move
                    Debug.Log($"Unit {unitAtCell.UnitId} must skip move (defense speed nerf)");
                    SelectUnit(null);
                    return;
                }

                SelectUnit(unitAtCell);
            }
            else
            {
                // Clicked on empty cell or enemy - deselect
                SelectUnit(null);
            }
        }

        // Right click to deselect
        if (mouse.rightButton.wasPressedThisFrame)
        {
            SelectUnit(null);
        }
    }

    void SelectUnit(GridUnit unit)
    {
        selectedUnit = unit;
        ClearValidMoveHighlights();

        if (unit == null)
        {
            if (selectionHighlight != null)
                selectionHighlight.SetActive(false);
            validMoves.Clear();
            return;
        }

        // Show selection highlight
        if (selectionHighlight == null)
        {
            selectionHighlight = CreateSprite("SelectionHighlight", SelectedColor, CellSize * 0.95f, CellSize * 0.95f);
            selectionHighlight.GetComponent<SpriteRenderer>().sortingOrder = 1;
        }
        selectionHighlight.SetActive(true);
        selectionHighlight.transform.position = GridToWorld(unit.GridPosition);

        // Calculate and show valid moves
        validMoves = GetValidMoves(unit);
        foreach (var move in validMoves)
        {
            var highlight = CreateSprite("ValidMove", ValidMoveColor, CellSize * 0.9f, CellSize * 0.9f);
            highlight.GetComponent<SpriteRenderer>().sortingOrder = 0;
            highlight.transform.position = GridToWorld(move);
            validMoveHighlights.Add(highlight);
        }
    }

    void ClearValidMoveHighlights()
    {
        foreach (var highlight in validMoveHighlights)
        {
            Destroy(highlight);
        }
        validMoveHighlights.Clear();
    }

    public List<Vector2Int> GetValidMoves(GridUnit unit)
    {
        var moves = new List<Vector2Int>();
        if (unit.IsTaggedOut) return moves;

        var pos = unit.GridPosition;
        bool isDefending = IsDefending(unit);

        // Determine valid directions based on zone
        Vector2Int[] directions;
        if (isDefending)
        {
            // Diagonal movement (8 directions) in own defense zone
            directions = new Vector2Int[]
            {
                new(-1, -1), new(0, -1), new(1, -1),
                new(-1, 0),              new(1, 0),
                new(-1, 1),  new(0, 1),  new(1, 1)
            };
        }
        else
        {
            // Orthogonal movement (4 directions) in offense/neutral zones
            directions = new Vector2Int[]
            {
                new(0, -1), new(-1, 0), new(1, 0), new(0, 1)
            };
        }

        foreach (var dir in directions)
        {
            var newPos = pos + dir;
            if (IsInBounds(newPos))
            {
                // Can move to empty cells or cells with enemies (will trigger collision)
                var occupant = cellOccupants[newPos.x, newPos.y];
                if (occupant == null || occupant.Team != unit.Team)
                {
                    moves.Add(newPos);
                }
            }
        }

        return moves;
    }

    void ExecutePlayerMove(GridUnit unit, Vector2Int destination)
    {
        ExecuteMove(unit, destination);
        EndTurn();
    }

    void ExecuteMove(GridUnit unit, Vector2Int destination)
    {
        var startPos = unit.GridPosition;
        bool wasDefending = IsDefending(unit);

        // Move unit
        MoveUnit(unit, destination);

        // Update defense move counter
        bool nowDefending = IsDefending(unit);
        if (nowDefending)
        {
            unit.ConsecutiveDefenseMoves++;
        }
        else
        {
            unit.ConsecutiveDefenseMoves = 0;
        }

        // Check collision
        CheckCollision(destination);

        // Check flag pickup
        CheckFlagPickup(unit);

        // Check scoring
        CheckScoring(unit);

        // Update visual
        unit.SetWorldPosition(GridToWorld(destination));
        UpdateFlagVisuals();
    }

    public void MoveUnit(GridUnit unit, Vector2Int to)
    {
        var from = unitPositions[unit];
        if (cellOccupants[from.x, from.y] == unit)
        {
            cellOccupants[from.x, from.y] = null;
        }
        cellOccupants[to.x, to.y] = unit;
        unitPositions[unit] = to;
        unit.GridPosition = to;
    }

    void CheckCollision(Vector2Int cell)
    {
        // Find all units at this cell
        var unitsHere = new List<GridUnit>();
        foreach (var unit in blueUnits)
        {
            if (!unit.IsTaggedOut && unit.GridPosition == cell)
                unitsHere.Add(unit);
        }
        foreach (var unit in redUnits)
        {
            if (!unit.IsTaggedOut && unit.GridPosition == cell)
                unitsHere.Add(unit);
        }

        if (unitsHere.Count < 2) return;

        // Check for enemy collisions
        for (int i = 0; i < unitsHere.Count; i++)
        {
            for (int j = i + 1; j < unitsHere.Count; j++)
            {
                var a = unitsHere[i];
                var b = unitsHere[j];
                if (a.Team != b.Team)
                {
                    ResolveCollision(a, b, cell);
                }
            }
        }
    }

    void ResolveCollision(GridUnit unitA, GridUnit unitB, Vector2Int cell)
    {
        if (unitA.IsTaggedOut || unitB.IsTaggedOut) return;

        var zoneOwner = GetZoneOwner(cell.y);

        // Defender in their zone always wins
        if (unitA.Team == Team.Blue && zoneOwner == ZoneOwner.Blue)
            TagOut(unitB);
        else if (unitB.Team == Team.Blue && zoneOwner == ZoneOwner.Blue)
            TagOut(unitA);
        else if (unitA.Team == Team.Red && zoneOwner == ZoneOwner.Red)
            TagOut(unitB);
        else if (unitB.Team == Team.Red && zoneOwner == ZoneOwner.Red)
            TagOut(unitA);
        else
        {
            // Neutral zone: lower UnitId wins (deterministic)
            if (unitA.UnitId < unitB.UnitId)
                TagOut(unitB);
            else
                TagOut(unitA);
        }
    }

    void TagOut(GridUnit unit)
    {
        if (unit.IsTaggedOut) return;

        unit.IsTaggedOut = true;
        unit.RespawnTurnsRemaining = RespawnDelay;
        unit.ConsecutiveDefenseMoves = 0;

        // Drop flag if carrying
        if (blueFlagCarrier == unit)
        {
            blueFlagCarrier = null;
            // Flag stays at current position
            Debug.Log("Blue flag dropped!");
        }
        if (redFlagCarrier == unit)
        {
            redFlagCarrier = null;
            Debug.Log("Red flag dropped!");
        }

        // Remove from cell
        var pos = unit.GridPosition;
        if (cellOccupants[pos.x, pos.y] == unit)
        {
            cellOccupants[pos.x, pos.y] = null;
        }

        // Fade out visual
        if (unit.SpriteRenderer != null)
        {
            var color = unit.SpriteRenderer.color;
            color.a = 0.3f;
            unit.SpriteRenderer.color = color;
        }

        Debug.Log($"{unit.Team} unit {unit.UnitId} tagged out!");
    }

    void CheckFlagPickup(GridUnit unit)
    {
        if (unit.IsTaggedOut) return;

        var pos = unit.GridPosition;

        // Blue unit can pick up red flag
        if (unit.Team == Team.Blue && pos == redFlagPosition && redFlagCarrier == null)
        {
            redFlagCarrier = unit;
            Debug.Log($"Blue unit picked up red flag!");
        }
        // Red unit can pick up blue flag
        else if (unit.Team == Team.Red && pos == blueFlagPosition && blueFlagCarrier == null)
        {
            blueFlagCarrier = unit;
            Debug.Log($"Red unit picked up blue flag!");
        }

        // Return own flag if it's dropped and friendly unit touches it
        if (unit.Team == Team.Blue && pos == blueFlagPosition && blueFlagCarrier == null && blueFlagPosition != blueFlagHome)
        {
            blueFlagPosition = blueFlagHome;
            Debug.Log("Blue flag returned home!");
        }
        if (unit.Team == Team.Red && pos == redFlagPosition && redFlagCarrier == null && redFlagPosition != redFlagHome)
        {
            redFlagPosition = redFlagHome;
            Debug.Log("Red flag returned home!");
        }
    }

    void CheckScoring(GridUnit unit)
    {
        if (unit.IsTaggedOut) return;

        var pos = unit.GridPosition;
        var zone = GetZoneOwner(pos.y);

        // Blue scores by bringing red flag to blue zone
        if (unit.Team == Team.Blue && redFlagCarrier == unit && zone == ZoneOwner.Blue)
        {
            Score(Team.Blue);
        }
        // Red scores by bringing blue flag to red zone
        else if (unit.Team == Team.Red && blueFlagCarrier == unit && zone == ZoneOwner.Red)
        {
            Score(Team.Red);
        }
    }

    void Score(Team team)
    {
        if (team == Team.Blue)
        {
            blueScore++;
            Debug.Log($"Blue scores! {blueScore} - {redScore}");
        }
        else
        {
            redScore++;
            Debug.Log($"Red scores! {blueScore} - {redScore}");
        }

        UpdateUI();

        if (blueScore >= WinScore)
        {
            EndGame(Team.Blue);
        }
        else if (redScore >= WinScore)
        {
            EndGame(Team.Red);
        }
        else
        {
            ResetRound();
        }
    }

    void ResetRound()
    {
        // Return flags
        blueFlagPosition = blueFlagHome;
        redFlagPosition = redFlagHome;
        blueFlagCarrier = null;
        redFlagCarrier = null;

        // Respawn all units
        for (int i = 0; i < blueUnits.Count; i++)
        {
            var unit = blueUnits[i];
            int x = ZoneBoundaries.BoardWidth / 2 - 1 + (i % 3);
            int y = 2 + (i / 3);
            RespawnUnit(unit, new Vector2Int(x, y));
        }

        for (int i = 0; i < redUnits.Count; i++)
        {
            var unit = redUnits[i];
            int x = ZoneBoundaries.BoardWidth / 2 - 1 + (i % 3);
            int y = ZoneBoundaries.BoardHeight - 3 - (i / 3);
            RespawnUnit(unit, new Vector2Int(x, y));
        }

        UpdateFlagVisuals();
        RecalculateVisibility();
    }

    void RespawnUnit(GridUnit unit, Vector2Int position)
    {
        unit.IsTaggedOut = false;
        unit.RespawnTurnsRemaining = 0;
        unit.ConsecutiveDefenseMoves = 0;

        // Clear old position
        var oldPos = unit.GridPosition;
        if (cellOccupants[oldPos.x, oldPos.y] == unit)
        {
            cellOccupants[oldPos.x, oldPos.y] = null;
        }

        // Find nearest empty cell if position is occupied
        if (cellOccupants[position.x, position.y] != null)
        {
            position = FindNearestEmptyCell(position);
        }

        unit.GridPosition = position;
        unitPositions[unit] = position;
        cellOccupants[position.x, position.y] = unit;
        unit.SetWorldPosition(GridToWorld(position));

        // Restore visual
        if (unit.SpriteRenderer != null)
        {
            var color = unit.SpriteRenderer.color;
            color.a = 1f;
            unit.SpriteRenderer.color = color;
        }
    }

    Vector2Int FindNearestEmptyCell(Vector2Int center)
    {
        for (int radius = 1; radius < 10; radius++)
        {
            for (int dx = -radius; dx <= radius; dx++)
            {
                for (int dy = -radius; dy <= radius; dy++)
                {
                    var pos = center + new Vector2Int(dx, dy);
                    if (IsInBounds(pos) && cellOccupants[pos.x, pos.y] == null)
                    {
                        return pos;
                    }
                }
            }
        }
        return center;
    }

    void EndTurn()
    {
        SelectUnit(null);
        turnNumber++;

        // Process respawn timers
        ProcessRespawns();

        // Recalculate visibility
        RecalculateVisibility();

        // Switch teams
        currentTeam = currentTeam == Team.Blue ? Team.Red : Team.Blue;

        UpdateUI();

        Debug.Log($"Turn {turnNumber}: {currentTeam}'s turn");
    }

    void ProcessRespawns()
    {
        foreach (var unit in blueUnits)
        {
            if (unit.IsTaggedOut)
            {
                unit.RespawnTurnsRemaining--;
                if (unit.RespawnTurnsRemaining <= 0)
                {
                    int idx = blueUnits.IndexOf(unit);
                    int x = ZoneBoundaries.BoardWidth / 2 - 1 + (idx % 3);
                    int y = 2 + (idx / 3);
                    RespawnUnit(unit, new Vector2Int(x, y));
                    Debug.Log($"Blue unit {unit.UnitId} respawned!");
                }
            }
        }

        foreach (var unit in redUnits)
        {
            if (unit.IsTaggedOut)
            {
                unit.RespawnTurnsRemaining--;
                if (unit.RespawnTurnsRemaining <= 0)
                {
                    int idx = redUnits.IndexOf(unit);
                    int x = ZoneBoundaries.BoardWidth / 2 - 1 + (idx % 3);
                    int y = ZoneBoundaries.BoardHeight - 3 - (idx / 3);
                    RespawnUnit(unit, new Vector2Int(x, y));
                    Debug.Log($"Red unit {unit.UnitId} respawned!");
                }
            }
        }
    }

    void EndGame(Team winner)
    {
        gameOver = true;

        if (gameOverText != null)
        {
            gameOverText.text = winner == Team.Blue ? "BLUE WINS!" : "RED WINS!";
            gameOverText.color = winner == Team.Blue ? Color.blue : Color.red;
            gameOverText.gameObject.SetActive(true);
        }

        Debug.Log($"{winner} wins!");
    }

    void UpdateUI()
    {
        if (scoreText != null)
        {
            scoreText.text = $"{blueScore} - {redScore}";
        }

        if (turnText != null)
        {
            string turnIndicator = currentTeam == Team.Blue ? "BLUE'S TURN" : "RED'S TURN";
            turnText.text = turnIndicator;
            turnText.color = currentTeam == Team.Blue ? Color.blue : Color.red;
        }
    }

    void UpdateFlagVisuals()
    {
        // Blue flag follows carrier or stays at position
        if (blueFlagCarrier != null && !blueFlagCarrier.IsTaggedOut)
        {
            blueFlagPosition = blueFlagCarrier.GridPosition;
            blueFlagGO.transform.position = GridToWorld(blueFlagPosition) + new Vector2(0.2f, 0.2f);
        }
        else
        {
            blueFlagGO.transform.position = GridToWorld(blueFlagPosition);
        }

        // Red flag follows carrier or stays at position
        if (redFlagCarrier != null && !redFlagCarrier.IsTaggedOut)
        {
            redFlagPosition = redFlagCarrier.GridPosition;
            redFlagGO.transform.position = GridToWorld(redFlagPosition) + new Vector2(0.2f, 0.2f);
        }
        else
        {
            redFlagGO.transform.position = GridToWorld(redFlagPosition);
        }

        // Apply fog of war visibility to flags
        bool blueFlagVisible = visibleToBlue.Contains(blueFlagPosition) || blueFlagCarrier != null;
        bool redFlagVisible = visibleToBlue.Contains(redFlagPosition) || redFlagCarrier != null;

        // For now, in single-player, Blue is the human player, so we apply Blue's visibility
        redFlagGO.GetComponent<SpriteRenderer>().enabled = redFlagVisible;
        // Blue flag is always visible to blue player
        blueFlagGO.GetComponent<SpriteRenderer>().enabled = true;
    }

    // Fog of War
    void RecalculateVisibility()
    {
        visibleToBlue.Clear();
        visibleToRed.Clear();

        foreach (var unit in blueUnits)
        {
            if (unit.IsTaggedOut) continue;
            AddVisibleCells(unit.GridPosition, visibleToBlue);
        }

        foreach (var unit in redUnits)
        {
            if (unit.IsTaggedOut) continue;
            AddVisibleCells(unit.GridPosition, visibleToRed);
        }

        // Update enemy visibility based on Blue's vision (human player)
        foreach (var enemy in redUnits)
        {
            bool visible = visibleToBlue.Contains(enemy.GridPosition) || enemy.IsTaggedOut;
            if (enemy.SpriteRenderer != null)
            {
                enemy.SpriteRenderer.enabled = visible;
            }
        }

        UpdateFlagVisuals();
    }

    void AddVisibleCells(Vector2Int center, HashSet<Vector2Int> visibleSet)
    {
        for (int dx = -VisionRadius; dx <= VisionRadius; dx++)
        {
            for (int dy = -VisionRadius; dy <= VisionRadius; dy++)
            {
                var cell = center + new Vector2Int(dx, dy);
                if (IsInBounds(cell))
                {
                    visibleSet.Add(cell);
                }
            }
        }
    }

    // AI interface
    public void ExecuteAIMove(GridUnit unit, Vector2Int destination)
    {
        if (gameOver) return;
        if (unit.Team != currentTeam) return;

        // Check speed nerf
        if (IsDefending(unit) && unit.ConsecutiveDefenseMoves % 4 == 3)
        {
            Debug.Log($"AI unit {unit.UnitId} skips move (defense speed nerf)");
            EndTurn();
            return;
        }

        ExecuteMove(unit, destination);
        EndTurn();
    }

    public void AISkipTurn()
    {
        if (gameOver) return;
        EndTurn();
    }

    public Team GetCurrentTeam() => currentTeam;
    public bool IsGameOver() => gameOver;
    public List<GridUnit> GetUnits(Team team) => team == Team.Blue ? blueUnits : redUnits;
    public HashSet<Vector2Int> GetVisibleCells(Team team) => team == Team.Blue ? visibleToBlue : visibleToRed;
    public Vector2Int GetEnemyFlagPosition(Team team) => team == Team.Blue ? redFlagPosition : blueFlagPosition;
    public GridUnit GetFlagCarrier(Team flagTeam) => flagTeam == Team.Blue ? blueFlagCarrier : redFlagCarrier;

    // Utility functions
    public ZoneOwner GetZoneOwner(int y) => y switch
    {
        < ZoneBoundaries.TeamBlueDefenseEnd => ZoneOwner.Blue,
        < ZoneBoundaries.NeutralEnd => ZoneOwner.Neutral,
        _ => ZoneOwner.Red
    };

    public bool IsDefending(GridUnit unit)
    {
        var zone = GetZoneOwner(unit.GridPosition.y);
        return (unit.Team == Team.Blue && zone == ZoneOwner.Blue) ||
               (unit.Team == Team.Red && zone == ZoneOwner.Red);
    }

    public bool IsInBounds(Vector2Int pos) =>
        pos.x >= 0 && pos.x < ZoneBoundaries.BoardWidth &&
        pos.y >= 0 && pos.y < ZoneBoundaries.BoardHeight;

    public Vector2 GridToWorld(Vector2Int gridPos)
    {
        // Center the board
        float offsetX = -ZoneBoundaries.BoardWidth * CellSize / 2f + CellSize / 2f;
        float offsetY = -ZoneBoundaries.BoardHeight * CellSize / 2f + CellSize / 2f;
        return new Vector2(
            gridPos.x * CellSize + offsetX,
            gridPos.y * CellSize + offsetY
        );
    }

    public Vector2Int WorldToGrid(Vector2 worldPos)
    {
        float offsetX = -ZoneBoundaries.BoardWidth * CellSize / 2f;
        float offsetY = -ZoneBoundaries.BoardHeight * CellSize / 2f;
        int x = Mathf.FloorToInt((worldPos.x - offsetX) / CellSize);
        int y = Mathf.FloorToInt((worldPos.y - offsetY) / CellSize);
        return new Vector2Int(x, y);
    }

    // Helper to create a colored sprite
    GameObject CreateSprite(string name, Color color, float width, float height)
    {
        var go = new GameObject(name);
        var sr = go.AddComponent<SpriteRenderer>();
        sr.sprite = GetRectSprite();
        sr.color = color;
        go.transform.localScale = new Vector3(width, height, 1);
        return go;
    }

    static Sprite cachedSprite;
    static Sprite GetRectSprite()
    {
        if (cachedSprite != null) return cachedSprite;

        var tex = new Texture2D(1, 1);
        tex.SetPixel(0, 0, Color.white);
        tex.Apply();

        cachedSprite = Sprite.Create(tex, new Rect(0, 0, 1, 1), new Vector2(0.5f, 0.5f), 1);
        return cachedSprite;
    }
}