from __future__ import annotations

import numpy as np

from blokus_gym.core.bots import GreedyBot, GreedyCornersBot, RandomBot
from blokus_gym.core.game import BlokusGame


class TestBlokusGame:
    def test_game_creation(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        assert game.board_size == 20
        assert game.num_players == 4
        assert game.current_player == 0
        assert len(game.players) == 4

    def test_reset(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        assert game.current_player == 0
        for player in game.players:
            assert len(player.available_pieces) == 21

    def test_action_space_size(self):
        game = BlokusGame(board_size=20, num_players=4)
        assert game.num_actions > 0

    def test_valid_actions_initial(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        valid = game.get_valid_actions(0)
        assert np.any(valid)

    def test_play_first_move_corner(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        valid = game.get_valid_actions(0)
        valid_indices = np.where(valid)[0]
        action = valid_indices[0]
        assert game.play_move(0, action)
        assert game.board.grid.sum() > 0

    def test_invalid_action(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        # Use an action that's likely invalid (very high index)
        assert not game.play_move(0, game.num_actions - 1)

    def test_next_player(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        assert game.current_player == 0
        game.next_player()
        assert game.current_player == 1
        game.next_player()
        assert game.current_player == 2
        game.next_player()
        assert game.current_player == 3
        game.next_player()
        assert game.current_player == 0

    def test_has_valid_moves(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        assert game.has_valid_moves(0)

    def test_is_game_over(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        assert not game.is_game_over()

    def test_get_scores(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        scores = game.get_scores()
        assert len(scores) == 4
        for score in scores:
            assert isinstance(score, (int, float))

    def test_get_action_mask(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        mask = game.get_action_mask(0)
        assert mask.shape == (game.num_actions,)
        assert mask.dtype == bool

    def test_get_current_observation(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        obs = game.get_current_observation()
        assert "board" in obs
        assert "pieces" in obs
        assert "corners" in obs
        assert obs["board"].shape == (20, 20)
        assert obs["pieces"].shape == (21,)
        assert obs["corners"].shape == (20, 20)

    def test_two_player_game(self):
        game = BlokusGame(board_size=14, num_players=2)
        game.reset()
        assert game.num_players == 2
        assert game.has_valid_moves(0)

    def test_three_player_game(self):
        game = BlokusGame(board_size=20, num_players=3)
        game.reset()
        assert game.num_players == 3

    def test_copy(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        game.play_move(0, 0)
        game_copy = game.copy()
        assert game_copy.current_player == game.current_player

    def test_play_multiple_moves(self):
        game = BlokusGame(board_size=20, num_players=2)
        game.reset()
        # Play a few moves
        for _ in range(5):
            valid = game.get_valid_actions(game.current_player)
            valid_indices = np.where(valid)[0]
            if len(valid_indices) > 0:
                game.play_move(game.current_player, valid_indices[0])
                game.next_player()
            else:
                break

    def test_get_winners_not_over(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        assert game.get_winners() is None

    def test_valid_move_first_corner_only(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        from blokus_gym.core.pieces import Move
        # First move must be in corner (0, 0) for player 0
        move = Move(piece_id=0, orientation_id=0, x=0, y=0)
        assert game.valid_move(0, move)
        # Try a non-corner first move
        move2 = Move(piece_id=0, orientation_id=0, x=5, y=5)
        assert not game.valid_move(0, move2)


class TestBots:
    def test_random_bot(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        bot = RandomBot(player_idx=0, seed=42)
        valid = game.get_valid_actions(0)
        action = bot.select_action(game, valid)
        assert action is not None
        assert valid[action]

    def test_greedy_bot(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        bot = GreedyBot(player_idx=0)
        valid = game.get_valid_actions(0)
        action = bot.select_action(game, valid)
        assert action is not None
        assert valid[action]

    def test_greedy_corners_bot(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        bot = GreedyCornersBot(player_idx=0)
        valid = game.get_valid_actions(0)
        action = bot.select_action(game, valid)
        assert action is not None
        assert valid[action]

    def test_bot_no_valid_moves(self):
        game = BlokusGame(board_size=20, num_players=4)
        game.reset()
        bot = RandomBot(player_idx=0, seed=42)
        valid = np.zeros(game.num_actions, dtype=bool)
        action = bot.select_action(game, valid)
        assert action is None

    def test_bots_play_game(self):
        game = BlokusGame(board_size=7, num_players=2)
        game.reset()
        bots = [RandomBot(player_idx=i, seed=i) for i in range(2)]
        steps = 0
        while not game.is_game_over() and steps < 100:
            bot = bots[game.current_player]
            valid = game.get_valid_actions(game.current_player)
            action = bot.select_action(game, valid)
            if action is not None:
                game.play_move(game.current_player, action)
            game.next_player()
            steps += 1
        assert game.is_game_over()