from __future__ import annotations

from typing import Any

import gymnasium as gym
import numpy as np
from gymnasium import spaces

from blokus_gym.core.bots import Bot, RandomBot
from blokus_gym.core.game import BlokusGame
from blokus_gym.core.pieces import STANDARD_PIECES, Piece, PieceSet
from blokus_gym.utils.render import render_ansi_board, render_text_board


class BlokusEnv(gym.Env):
    """Gymnasium environment for the board game Blokus.

    Supports 2-4 players on customizable board sizes with custom piece sets.
    In single-agent mode, the agent controls player 0 and opponents are bots.

    Action Space: Discrete(N) where N is the total number of possible
    (piece, orientation, position) combinations. An action mask is provided
    in the info dict to indicate valid actions.

    Observation Space: Dict with:
        - "board": (board_size, board_size) int8 array (0=empty, 1-4=player)
        - "pieces": MultiBinary(num_pieces) - 1 if piece still available
        - "corners": (board_size, board_size) bool array - valid corner cells

    Example:
        >>> import gymnasium as gym
        >>> from blokus_gym import BlokusEnv
        >>> env = BlokusEnv(num_players=4, board_size=20)
        >>> obs, info = env.reset(seed=42)
        >>> action = env.action_space.sample()  # Use info["action_mask"] instead
        >>> obs, reward, terminated, truncated, info = env.step(action)
    """

    metadata = {"render_modes": ["human", "ansi", "rgb_array"]}

    # Player colors for rendering (index 0 = empty)
    PLAYER_COLORS = ["empty", "red", "blue", "yellow", "green"]

    def __init__(
        self,
        num_players: int = 4,
        board_size: int = 20,
        pieces: list[Piece] | None = None,
        render_mode: str | None = None,
        bot_type: type[Bot] = RandomBot,
        bot_strength: int = 1,
        reward_shaping: bool = False,
        corner_rule: bool = True,
        max_steps: int | None = None,
        seed: int | None = None,
        **kwargs: Any,
    ):
        super().__init__()

        # Validate parameters
        assert 2 <= num_players <= 4, f"num_players must be 2-4, got {num_players}"
        assert board_size >= 5, f"board_size must be >= 5, got {board_size}"

        self.num_players = num_players
        self.board_size = board_size
        self.render_mode = render_mode
        self.bot_type = bot_type
        self.bot_strength = bot_strength
        self.reward_shaping = reward_shaping
        self.corner_rule = corner_rule
        self.max_steps = max_steps
        self._seed = seed

        # Initialize piece set
        self.piece_set = PieceSet(pieces or STANDARD_PIECES)

        # Initialize game
        self.game = BlokusGame(
            board_size=board_size,
            pieces=self.piece_set,
            num_players=num_players,
            corner_rule=corner_rule,
        )

        # Define spaces
        self.observation_space = spaces.Dict({
            "board": spaces.Box(0, num_players, (board_size, board_size), dtype=np.int8),
            "pieces": spaces.MultiBinary(self.piece_set.num_pieces),
            "corners": spaces.Box(0, 1, (board_size, board_size), dtype=bool),
        })

        self.action_space = spaces.Discrete(self.game.num_actions)

        # Bot instances (created during reset)
        self.bots: list[Bot | None] = [None] * num_players

        # Episode tracking
        self.current_step = 0
        self._last_observation: dict | None = None

    def _get_obs(self) -> dict:
        """Get the current observation from the agent's perspective (player 0)."""
        obs = self.game.get_current_observation()
        return {
            "board": obs["board"],
            "pieces": obs["pieces"],
            "corners": obs["corners"],
        }

    def _get_info(self) -> dict:
        """Get auxiliary information."""
        return {
            "action_mask": self.game.get_action_mask(0),
            "current_player": self.game.current_player,
            "players_with_moves": [
                i for i in range(self.num_players)
                if self.game.players[i].can_move and self.game.has_valid_moves(i)
            ],
            "step_count": self.current_step,
            "scores": self.game.get_scores(),
        }

    def reset(
        self,
        seed: int | None = None,
        options: dict | None = None,
    ) -> tuple[dict, dict]:
        """Reset the environment to a new episode.

        Args:
            seed: Random seed for reproducibility.
            options: Additional options (unused).

        Returns:
            Tuple of (observation, info).
        """
        super().reset(seed=seed)

        # Re-initialize game
        self.game = BlokusGame(
            board_size=self.board_size,
            pieces=self.piece_set,
            num_players=self.num_players,
            corner_rule=self.corner_rule,
        )
        self.game.reset()

        # Create bots for opponents
        bot_seed = seed or 0
        for i in range(1, self.num_players):
            self.bots[i] = self.bot_type(
                player_idx=i,
                seed=bot_seed + i * 1000,
            )

        self.current_step = 0

        # Let bots play until it's the agent's turn
        # The agent (player 0) starts first if corner_rule is enabled
        self._play_bots_until_agent()

        obs = self._get_obs()
        info = self._get_info()
        self._last_observation = obs

        return obs, info

    def _play_bots_until_agent(self) -> None:
        """Let bot players take their turns until it's the agent's turn."""
        max_bot_turns = self.num_players * 10  # Safety limit
        turns = 0

        while self.game.current_player != 0 and not self.game.is_game_over():
            if turns > max_bot_turns:
                break
            bot = self.bots[self.game.current_player]
            if bot is None:
                # Shouldn't happen, but safety
                self.game.next_player()
                turns += 1
                continue

            valid_actions = self.game.get_valid_actions(self.game.current_player)
            if not np.any(valid_actions):
                self.game.players[self.game.current_player].can_move = False
                self.game.next_player()
                turns += 1
                continue

            action = bot.select_action(self.game, valid_actions)
            if action is not None:
                self.game.play_move(self.game.current_player, action)
            else:
                self.game.players[self.game.current_player].can_move = False

            self.game.next_player()
            turns += 1

    def step(self, action: int) -> tuple[dict, float, bool, bool, dict]:
        """Execute one step in the environment.

        The agent (player 0) takes an action, then all bots play until
        it's the agent's turn again.

        Args:
            action: Action index (piece, orientation, position).

        Returns:
            Tuple of (observation, reward, terminated, truncated, info).
        """
        self.current_step += 1

        # Apply agent's action
        success = self.game.play_move(0, action)

        if not success:
            # Invalid action - penalize and end episode
            reward = -10.0
            obs = self._get_obs()
            info = self._get_info()
            info["invalid_action"] = True
            self._last_observation = obs
            return obs, reward, True, False, info

        # Advance to next player
        self.game.next_player()

        # Let bots play until it's the agent's turn or game is over
        self._play_bots_until_agent()

        # Check if game is over
        terminated = self.game.is_game_over()
        truncated = False
        if self.max_steps is not None and self.current_step >= self.max_steps:
            truncated = True

        # Compute reward
        if terminated or truncated:
            reward = self._compute_terminal_reward()
        else:
            reward = self._compute_step_reward()

        obs = self._get_obs()
        info = self._get_info()
        self._last_observation = obs

        return obs, reward, terminated, truncated, info

    def _compute_step_reward(self) -> float:
        """Compute reward for a non-terminal step."""
        if self.reward_shaping:
            # Small reward for each square placed
            return 0.01 * self.game.players[0].score
        return 0.0

    def _compute_terminal_reward(self) -> float:
        """Compute reward when the game ends."""
        scores = self.game.get_scores()
        agent_score = scores[0]

        if self.reward_shaping:
            # Normalized score: agent score / max possible score
            max_possible = sum(p.size for p in self.piece_set.pieces)
            return agent_score / max_possible

        # Win/loss reward
        max_score = max(scores)
        if agent_score == max_score:
            # Check if it's a tie
            winners = [i for i, s in enumerate(scores) if s == max_score]
            if len(winners) == 1:
                return 1.0  # Win
            else:
                return 0.0  # Tie
        else:
            return -1.0  # Loss

    def render(self) -> str | np.ndarray | None:
        """Render the environment.

        Args:
            No arguments - uses self.render_mode.

        Returns:
            For "human": None (prints to stdout)
            For "ansi": string representation
            For "rgb_array": numpy array of shape (H, W, 3)
        """
        if self.render_mode is None:
            return None

        if self.render_mode == "human":
            render_text_board(self.game)
            return None

        elif self.render_mode == "ansi":
            return render_ansi_board(self.game)

        elif self.render_mode == "rgb_array":
            return self._render_rgb_array()

        else:
            raise ValueError(f"Unknown render_mode: {self.render_mode}")

    def _render_rgb_array(self) -> np.ndarray:
        """Render the board as an RGB array using matplotlib."""
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError as exc:
            raise ImportError(
                "matplotlib is required for rgb_array rendering. "
                "Install with: pip install matplotlib"
            ) from exc

        fig, ax = plt.subplots(figsize=(8, 8))

        # Color map: 0=empty(light grey), 1=red, 2=blue, 3=yellow, 4=green
        colors = {
            0: "#d3d3d3",  # light grey
            1: "#ff6b6b",  # red
            2: "#4dabf7",  # blue
            3: "#ffd43b",  # yellow
            4: "#51cf66",  # green
        }

        grid = self.game.board.grid.astype(float)
        colored = np.zeros((self.board_size, self.board_size, 3))
        for i in range(self.board_size):
            for j in range(self.board_size):
                val = int(grid[i, j])
                hex_color = colors.get(val, "#d3d3d3")
                # Parse hex color
                r = int(hex_color[1:3], 16) / 255
                g = int(hex_color[3:5], 16) / 255
                b = int(hex_color[5:7], 16) / 255
                colored[i, j] = [r, g, b]

        ax.imshow(colored, interpolation="nearest")
        ax.set_xticks(np.arange(-0.5, self.board_size, 1), minor=True)
        ax.set_yticks(np.arange(-0.5, self.board_size, 1), minor=True)
        ax.grid(True, which="minor", color="black", linewidth=0.5)
        ax.set_xticks([])
        ax.set_yticks([])

        # Title with current player and scores
        scores = self.game.get_scores()
        title = f"Current player: {self.game.current_player + 1}\n"
        title += "Scores: " + " | ".join(f"P{i+1}: {s}" for i, s in enumerate(scores))
        ax.set_title(title, fontsize=10)

        plt.tight_layout()
        fig.canvas.draw()
        image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        plt.close(fig)
        return image

    def close(self) -> None:
        """Clean up resources."""
        pass

    def get_action_mask(self) -> np.ndarray:
        """Get the action mask for the current player (player 0)."""
        return self.game.get_action_mask(0)
