from __future__ import annotations

import random
from abc import ABC, abstractmethod

import numpy as np

from blokus_gym.core.game import BlokusGame


class Bot(ABC):
    """Base class for Blokus bots."""

    def __init__(self, player_idx: int, seed: int | None = None):
        self.player_idx = player_idx
        self.rng = random.Random(seed)

    @abstractmethod
    def select_action(self, game: BlokusGame, valid_actions: np.ndarray) -> int | None:
        """Select an action from the valid actions.

        Args:
            game: The current game state.
            valid_actions: Boolean mask of valid actions.

        Returns:
            Action index, or None if no valid moves.
        """
        ...

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(player={self.player_idx})"


class RandomBot(Bot):
    """Randomly selects from valid actions."""

    def select_action(self, game: BlokusGame, valid_actions: np.ndarray) -> int | None:
        valid_indices = np.where(valid_actions)[0]
        if len(valid_indices) == 0:
            return None
        return int(self.rng.choice(valid_indices))


class GreedyBot(Bot):
    """Selects the move that places the most squares (largest piece)."""

    def select_action(self, game: BlokusGame, valid_actions: np.ndarray) -> int | None:
        valid_indices = np.where(valid_actions)[0]
        if len(valid_indices) == 0:
            return None

        best_action = -1
        best_size = -1
        for action in valid_indices:
            move = game.get_move(int(action))
            piece = game.piece_set.get_piece(move.piece_id)
            if piece.size > best_size:
                best_size = piece.size
                best_action = int(action)

        return best_action


class GreedyCornersBot(Bot):
    """Greedy bot that prefers moves opening more corners for future play."""

    def select_action(self, game: BlokusGame, valid_actions: np.ndarray) -> int | None:
        valid_indices = np.where(valid_actions)[0]
        if len(valid_indices) == 0:
            return None

        best_action = -1
        best_score = float("-inf")

        for action in valid_indices:
            move = game.get_move(int(action))
            piece = game.piece_set.get_piece(move.piece_id)
            placed_corners = game._get_placed_corners(move)

            # Count how many new corners this move opens
            new_corners = 0
            for cx, cy in placed_corners:
                if game.board.in_bounds(cx, cy) and game.board.is_empty(cx, cy):
                    new_corners += 1

            # Score = piece size + corner bonus
            score = piece.size + new_corners * 0.5
            if score > best_score:
                best_score = score
                best_action = int(action)

        return best_action


class MinimaxBot(Bot):
    """Minimax search bot (primarily for 2-player games).

    Uses a depth-limited minimax search with a simple heuristic.
    """

    def __init__(self, player_idx: int, depth: int = 2, seed: int | None = None):
        super().__init__(player_idx, seed)
        self.depth = depth

    def select_action(self, game: BlokusGame, valid_actions: np.ndarray) -> int | None:
        valid_indices = np.where(valid_actions)[0]
        if len(valid_indices) == 0:
            return None

        if len(valid_indices) == 1:
            return int(valid_indices[0])

        best_action = -1
        best_value = float("-inf")

        for action in valid_indices:
            if game.num_players != 2:
                # For multi-player, just use greedy
                return GreedyBot(self.player_idx, self.rng.randint(0, 2**31)).select_action(
                    game, valid_actions
                )

            cloned = game.copy()
            success = cloned.play_move(self.player_idx, int(action))
            if not success:
                continue

            value = self._minimax(cloned, self.depth - 1, False)
            if value > best_value:
                best_value = value
                best_action = int(action)

        return best_action if best_action >= 0 else int(valid_indices[0])

    def _minimax(self, game: BlokusGame, depth: int, maximizing: bool) -> float:
        if depth == 0 or game.is_game_over():
            return self._evaluate(game)

        if maximizing:
            value = float("-inf")
            valid = game.get_valid_actions(self.player_idx)
            valid_indices = np.where(valid)[0]
            for action in valid_indices:
                cloned = game.copy()
                if cloned.play_move(self.player_idx, int(action)):
                    val = self._minimax(cloned, depth - 1, False)
                    value = max(value, val)
            return value
        else:
            # Opponent's turn
            opp_idx = 1 - self.player_idx
            value = float("inf")
            valid = game.get_valid_actions(opp_idx)
            valid_indices = np.where(valid)[0]
            for action in valid_indices:
                cloned = game.copy()
                if cloned.play_move(opp_idx, int(action)):
                    val = self._minimax(cloned, depth - 1, True)
                    value = min(value, val)
            return value

    def _evaluate(self, game: BlokusGame) -> float:
        """Simple heuristic: difference in placed squares."""
        my_score = game.players[self.player_idx].score
        opp_idx = 1 - self.player_idx
        opp_score = game.players[opp_idx].score if opp_idx < len(game.players) else 0
        return float(my_score - opp_score)
