from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

from blokus_gym.core.board import Board
from blokus_gym.core.pieces import STANDARD_PIECES, Move, PieceSet


@dataclass
class PlayerState:
    """Tracks the state of a single player in the game."""

    idx: int  # 0-based player index
    available_pieces: set[str] = field(default_factory=set)  # Piece names still available
    score: int = 0  # Squares placed (positive = good)
    corners: set[tuple[int, int]] = field(default_factory=set)  # Valid placement corners
    has_started: bool = False  # Whether player has placed their first piece
    can_move: bool = True  # Whether player can still make any move


class BlokusGame:
    """Core Blokus game logic.

    Manages the board state, player turns, move validation, and scoring.
    This class is independent of the Gymnasium environment and can be
    used standalone for game simulation or bot development.
    """

    # Player indices are 1-based for board representation
    # (0 = empty, 1 = player 1, etc.)
    PLAYER_OFFSET = 1

    def __init__(
        self,
        board_size: int = 20,
        pieces: PieceSet | None = None,
        num_players: int = 4,
        corner_rule: bool = True,
    ):
        self.board_size = board_size
        self.piece_set = pieces or PieceSet(STANDARD_PIECES)
        self.num_players = num_players
        self.corner_rule = corner_rule

        self.board = Board(board_size)
        self.players: list[PlayerState] = []
        self.current_player = 0
        self.rounds = 0
        self.game_over = False

        # Pre-compute all possible moves (action lookup table)
        self._action_moves: list[Move] = []
        self._move_to_action: dict[tuple[int, int, int, int], int] = {}
        self._generate_action_space()

        # Pre-compute action indices per piece for faster filtering
        self._actions_by_piece: list[list[int]] = []
        for piece_id in range(self.piece_set.num_pieces):
            self._actions_by_piece.append([
                idx for idx, move in enumerate(self._action_moves)
                if move.piece_id == piece_id
            ])

        # Starting corners for each player (standard Blokus layout)
        # 4 players: all four corners
        # 2 players: opposite corners (0,0) and (max, max)
        # 3 players: three corners
        max_idx = board_size - 1
        self._starting_corners = [
            (0, 0),
            (0, max_idx),
            (max_idx, 0),
            (max_idx, max_idx),
        ]
        if num_players == 2:
            self._starting_corners = [(0, 0), (max_idx, max_idx)]
        elif num_players == 3:
            self._starting_corners = [(0, 0), (0, max_idx), (max_idx, 0)]

    # ------------------------------------------------------------------
    # Action space generation
    # ------------------------------------------------------------------

    def _generate_action_space(self) -> None:
        """Pre-compute all possible (piece, orientation, position) moves.

        Each unique combination is assigned a stable integer action index.
        """
        self._action_moves = []
        self._move_to_action = {}

        for piece_id in range(self.piece_set.num_pieces):
            for orient_id, orient in enumerate(self.piece_set.get_orientations(piece_id)):
                squares = orient.squares
                # Compute bounding box of the piece
                max_x = max(x for x, _ in squares)
                max_y = max(y for _, y in squares)

                for px in range(self.board_size - max_x):
                    for py in range(self.board_size - max_y):
                        move = Move(
                            piece_id=piece_id,
                            orientation_id=orient_id,
                            x=px,
                            y=py,
                        )
                        action_idx = len(self._action_moves)
                        self._action_moves.append(move)
                        self._move_to_action[(piece_id, orient_id, px, py)] = action_idx

    @property
    def num_actions(self) -> int:
        return len(self._action_moves)

    def get_move(self, action: int) -> Move:
        """Get the Move object for a given action index."""
        return self._action_moves[action]

    def get_action(self, move: Move) -> int:
        """Get the action index for a given Move."""
        return self._move_to_action[(move.piece_id, move.orientation_id, move.x, move.y)]

    # ------------------------------------------------------------------
    # Game initialization
    # ------------------------------------------------------------------

    def reset(self) -> None:
        """Reset the game to its initial state."""
        self.board = Board(self.board_size)
        self.players = []
        self.current_player = 0
        self.rounds = 0
        self.game_over = False

        for i in range(self.num_players):
            player = PlayerState(idx=i)
            player.available_pieces = set(self.piece_set.piece_names)
            if self.corner_rule and i < len(self._starting_corners):
                player.corners = {self._starting_corners[i]}
            else:
                player.corners = set()
            self.players.append(player)

    # ------------------------------------------------------------------
    # Move validation
    # ------------------------------------------------------------------

    def _get_placed_squares(self, move: Move) -> list[tuple[int, int]]:
        """Get the absolute board coordinates for a move."""
        orient = self.piece_set.get_orientations(move.piece_id)[move.orientation_id]
        return [(move.x + dx, move.y + dy) for dx, dy in orient.squares]

    def _get_placed_corners(self, move: Move) -> list[tuple[int, int]]:
        """Get the absolute corner coordinates for a move."""
        orient = self.piece_set.get_orientations(move.piece_id)[move.orientation_id]
        return [(move.x + dx, move.y + dy) for dx, dy in orient.corners]

    def valid_move(self, player_idx: int, move: Move) -> bool:
        """Check if a move is valid for the given player.

        Validation rules:
        1. Player must have the piece available
        2. All squares must be in bounds
        3. No overlap with existing pieces
        4. Corner rule: piece must touch same-color piece at a corner
        5. No edge adjacency with same-color pieces
        6. First move must be in a corner (if corner_rule is enabled)
        """
        player = self.players[player_idx]
        piece = self.piece_set.get_piece(move.piece_id)

        # Rule 1: Check if player has the piece
        if piece.name not in player.available_pieces:
            return False

        placed_squares = self._get_placed_squares(move)

        # Rule 2: Check bounds
        for x, y in placed_squares:
            if not self.board.in_bounds(x, y):
                return False

        # Rule 3: Check overlap
        if self.board.has_overlap(placed_squares):
            return False

        player_board_idx = player_idx + self.PLAYER_OFFSET

        # Rule 6: First move must be in a corner
        if self.corner_rule and not player.has_started:
            if len(self.players) <= len(self._starting_corners):
                start_corner = self._starting_corners[player_idx]
                if (move.x, move.y) != start_corner:
                    return False

        # Rule 5: No edge adjacency with same-color pieces
        for x, y in placed_squares:
            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nx, ny = x + dx, y + dy
                if self.board.in_bounds(nx, ny):
                    if self.board.get_cell(nx, ny) == player_board_idx:
                        return False

        # Rule 4: Corner rule (must touch same-color at a corner)
        if player.has_started:
            placed_corners = self._get_placed_corners(move)
            # Use cached occupied squares from board for fast lookup
            player_occupied = self.board.get_player_occupied(player_board_idx)
            touches_corner = any(
                self.board.in_bounds(cx, cy) and (cx, cy) in player_occupied
                for cx, cy in placed_corners
            )
            if not touches_corner:
                return False

        return True

    def get_valid_actions(self, player_idx: int) -> np.ndarray:
        """Get a boolean mask of valid actions for the given player."""
        mask = np.zeros(self.num_actions, dtype=bool)
        player = self.players[player_idx]

        if not player.can_move:
            return mask

        available_piece_ids = [
            self.piece_set.get_piece_id(name)
            for name in player.available_pieces
        ]

        if not player.has_started or not player.corners:
            for piece_id in available_piece_ids:
                for action_idx in self._actions_by_piece[piece_id]:
                    if self.valid_move(player_idx, self._action_moves[action_idx]):
                        mask[action_idx] = True
            return mask

        player_board_idx = player_idx + self.PLAYER_OFFSET
        
        # Pre-compute player's occupied squares ONCE (cached)
        player_occupied = self.board.get_player_occupied(player_board_idx)

        for piece_id in available_piece_ids:
            for action_idx in self._actions_by_piece[piece_id]:
                move = self._action_moves[action_idx]
                placed_corners = self._get_placed_corners(move)

                # Early exit: check corner touch using cached occupied squares
                touches_corner = any(
                    self.board.in_bounds(cx, cy) and (cx, cy) in player_occupied
                    for cx, cy in placed_corners
                )
                if not touches_corner:
                    continue  # Skip expensive overlap/edge checks

                placed_squares = self._get_placed_squares(move)

                # Check bounds
                if not all(self.board.in_bounds(x, y) for x, y in placed_squares):
                    continue

                # Check overlap
                if self.board.has_overlap(placed_squares):
                    continue

                # Check edge adjacency (no same-color edges)
                edge_invalid = False
                for x, y in placed_squares:
                    for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                        nx, ny = x + dx, y + dy
                        if self.board.in_bounds(nx, ny) and self.board.get_cell(nx, ny) == player_board_idx:
                            edge_invalid = True
                            break
                    if edge_invalid:
                        break

                if not edge_invalid:
                    mask[action_idx] = True

        return mask

    def get_valid_action_indices(self, player_idx: int) -> list[int]:
        """Get a list of valid action indices for the given player."""
        mask = self.get_valid_actions(player_idx)
        return [int(i) for i in np.where(mask)[0]]

    def has_valid_moves(self, player_idx: int) -> bool:
        """Check if a player has any valid moves."""
        player = self.players[player_idx]
        if not player.can_move:
            return False

        available_piece_ids = [
            self.piece_set.get_piece_id(name)
            for name in player.available_pieces
        ]

        for piece_id in available_piece_ids:
            for action_idx in self._actions_by_piece[piece_id]:
                if self.valid_move(player_idx, self._action_moves[action_idx]):
                    return True
        return False

    # ------------------------------------------------------------------
    # Move execution
    # ------------------------------------------------------------------

    def apply_move(self, player_idx: int, move: Move) -> None:
        """Apply a move to the game state. Does not validate."""
        player = self.players[player_idx]
        piece = self.piece_set.get_piece(move.piece_id)
        placed_squares = self._get_placed_squares(move)
        placed_corners = self._get_placed_corners(move)

        # Place on board
        player_board_idx = player_idx + self.PLAYER_OFFSET
        self.board.place(player_board_idx, placed_squares)

        # Update player state
        player.score += piece.size
        player.available_pieces.discard(piece.name)
        player.has_started = True

        # Update corners: remove covered squares, add new corners
        for cx, cy in placed_corners:
            if self.board.in_bounds(cx, cy) and self.board.is_empty(cx, cy):
                player.corners.add((cx, cy))

        # Remove corners that are now occupied
        player.corners = {
            (x, y) for x, y in player.corners if self.board.is_empty(x, y)
        }

    def play_move(self, player_idx: int, action: int) -> bool:
        """Validate and apply a move. Returns True if successful."""
        move = self.get_move(action)
        if not self.valid_move(player_idx, move):
            return False
        self.apply_move(player_idx, move)
        return True

    # ------------------------------------------------------------------
    # Turn management
    # ------------------------------------------------------------------

    def next_player(self) -> int:
        """Advance to the next player who can still move.

        If no players can move, the game is over.
        """
        if self.game_over:
            return self.current_player

        for offset in range(1, self.num_players + 1):
            next_idx = (self.current_player + offset) % self.num_players
            next_player = self.players[next_idx]
            if next_player.can_move and self.has_valid_moves(next_idx):
                self.current_player = next_idx
                return next_idx

        # No one can move — game over
        self.game_over = True
        return self.current_player

    def advance_turn(self) -> None:
        """Apply a move for the current player, then advance.

        This is a convenience method for bot-driven games.
        """
        self.next_player()

    # ------------------------------------------------------------------
    # Game state queries
    # ------------------------------------------------------------------

    def is_game_over(self) -> bool:
        """Check if the game is over (no player can make a move)."""
        if self.game_over:
            return True
        # Check if any player can still move
        for i in range(self.num_players):
            if self.players[i].can_move and self.has_valid_moves(i):
                return False
        self.game_over = True
        return True

    def get_scores(self) -> list[int]:
        """Get the score for each player.

        Score = squares placed (positive) - unplaced squares (negative).
        A player who placed all pieces gets a +15 bonus.
        If the last piece was the monomino, an additional +5 bonus.
        """
        scores = []
        for player in self.players:
            # Sum sizes of remaining pieces
            unplaced = sum(
                self.piece_set.pieces[pid].size
                for pid, piece in enumerate(self.piece_set.pieces)
                if piece.name in player.available_pieces
            )
            score = player.score - unplaced

            # Bonus for placing all pieces
            if len(player.available_pieces) == 0:
                score += 15

            scores.append(score)
        return scores

    def get_winners(self) -> list[int] | None:
        """Get the winning player(s). Returns None if game is not over."""
        if not self.is_game_over():
            return None

        scores = self.get_scores()
        max_score = max(scores)
        winners = [i for i, s in enumerate(scores) if s == max_score]
        return winners

    def get_current_observation(self) -> dict:
        """Get the current game state as an observation dict."""
        player_idx = self.current_player
        player = self.players[player_idx]

        # Board as player indices (1-based)
        board_obs = self.board.grid.copy()

        # Available pieces as binary vector
        pieces_obs = np.zeros(self.piece_set.num_pieces, dtype=bool)
        for name in player.available_pieces:
            pieces_obs[self.piece_set.get_piece_id(name)] = True

        # Corners as boolean grid
        corners_obs = np.zeros((self.board_size, self.board_size), dtype=bool)
        for x, y in player.corners:
            if self.board.in_bounds(x, y):
                corners_obs[y, x] = True

        return {
            "board": board_obs.astype(np.int8),
            "pieces": pieces_obs,
            "corners": corners_obs,
        }

    def get_action_mask(self, player_idx: int) -> np.ndarray:
        """Get the action mask for a player."""
        return self.get_valid_actions(player_idx)

    def copy(self) -> BlokusGame:
        """Create a deep copy of the game state."""
        import copy

        new_game = BlokusGame.__new__(BlokusGame)
        new_game.board_size = self.board_size
        new_game.piece_set = self.piece_set
        new_game.num_players = self.num_players
        new_game.corner_rule = self.corner_rule
        new_game.board = self.board.copy()
        new_game.players = copy.deepcopy(self.players)
        new_game.current_player = self.current_player
        new_game.rounds = self.rounds
        new_game.game_over = self.game_over
        new_game._action_moves = self._action_moves
        new_game._move_to_action = self._move_to_action
        new_game._starting_corners = self._starting_corners
        return new_game
