from __future__ import annotations

import numpy as np
from gymnasium import spaces

from blokus_gym.core.game import BlokusGame
from blokus_gym.core.pieces import STANDARD_PIECES, PieceSet


class BlokusMultiAgentEnv:
    """PettingZoo-style AEC (Actor-Environment-Cycle) environment for Blokus.

    This wraps the BlokusGame core to provide a multi-agent interface where
    each player acts independently. Compatible with PettingZoo's API and
    RLlib's multi-agent RL.

    Unlike the single-agent BlokusEnv, all players are controlled by the
    caller (no built-in bots). This allows for self-play training and
    evaluation of multi-agent policies.

    Attributes:
        num_players: Number of players (2-4).
        board_size: Size of the square board.
        agents: List of agent names (e.g., ["player_0", "player_1", ...]).
        possible_agents: Same as agents (never changes).
        agent_order: Order in which agents take turns.

    Example:
        >>> from blokus_gym import BlokusMultiAgentEnv
        >>> env = BlokusMultiAgentEnv(num_players=4)
        >>> obs, info = env.reset()
        >>> for agent in env.agent_iter():
        ...     obs, reward, terminated, truncated, info = env.last()
        ...     action = env.action_space(agent).sample()  # Use action mask
        ...     env.step(action)
    """

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

    def __init__(
        self,
        num_players: int = 4,
        board_size: int = 20,
        pieces: list | None = None,
        render_mode: str | None = None,
        corner_rule: bool = True,
        seed: int | None = None,
    ):
        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.corner_rule = corner_rule
        self._seed = seed

        self.piece_set = PieceSet(pieces or STANDARD_PIECES)
        self.game: BlokusGame | None = None

        self.agents = [f"player_{i}" for i in range(num_players)]
        self.possible_agents = list(self.agents)
        self.agent_order = list(self.agents)

        # Observation and action spaces (same for all agents)
        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._cur_step = 0
        self._current_agent_idx = 0
        self._agent_dones: dict[str, bool] = {}

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

        Returns:
            obs: Dict mapping agent name to observation.
            info: Dict with auxiliary information.
        """
        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()

        self._cur_step = 0
        self._current_agent_idx = 0
        self._agent_dones = dict.fromkeys(self.agents, False)

        # Get observation for the first agent
        first_agent = self.agent_order[0]
        obs = self._get_obs(first_agent)
        info = self._get_info(first_agent)

        return {first_agent: obs}, info

    def step(self, action: int) -> tuple[dict, dict, dict, dict, dict]:
        """Execute one step. Must be called for the current agent in agent_iter().

        Returns:
            observations: Dict mapping agent -> observation
            rewards: Dict mapping agent -> reward
            terminations: Dict mapping agent -> bool
            truncations: Dict mapping agent -> bool
            infos: Dict mapping agent -> info
        """
        self._cur_step += 1
        agent = self.agent_order[self._current_agent_idx]
        player_idx = int(agent.split("_")[1])

        # Apply the action
        success = self.game.play_move(player_idx, action)

        # Compute rewards for this player
        reward = 0.0
        if not success:
            reward = -10.0  # Invalid action penalty
            self._agent_dones[agent] = True

        # Check game over
        game_over = self.game.is_game_over()

        # Advance to next player
        if not game_over:
            self.game.next_player()
            # Skip players who can't move
            while not self.game.is_game_over():
                player = self.game.current_player
                agent_name = self.agents[player]
                if self._agent_dones.get(agent_name, False):
                    self.game.next_player()
                    continue
                if not self.game.has_valid_moves(player):
                    self.game.players[player].can_move = False
                    self.game.next_player()
                    continue
                break

        # Determine next agent
        next_player = self.game.current_player
        next_agent = self.agents[next_player]

        # Check if game is over
        if game_over:
            # Final rewards based on scores
            scores = self.game.get_scores()
            max_score = max(scores)
            terminations = {}
            rewards = {}
            for i, agent_name in enumerate(self.agents):
                if scores[i] == max_score:
                    winners = [j for j, s in enumerate(scores) if s == max_score]
                    if len(winners) == 1:
                        rewards[agent_name] = 1.0
                    else:
                        rewards[agent_name] = 0.0
                else:
                    rewards[agent_name] = -1.0
                terminations[agent_name] = True
                self._agent_dones[agent_name] = True

            infos = {agent_name: {"step_count": self._cur_step} for agent_name in self.agents}
            observations = {agent_name: self._get_obs(agent_name) for agent_name in self.agents}

            return observations, rewards, terminations, dict.fromkeys(self.agents, False), infos

        # Normal step
        rewards = dict.fromkeys(self.agents, 0.0)
        rewards[agent] = reward
        terminations = {
            agent_name: self._agent_dones.get(agent_name, False)
            for agent_name in self.agents
        }
        truncations = dict.fromkeys(self.agents, False)

        observations = {next_agent: self._get_obs(next_agent)}
        infos = {next_agent: self._get_info(next_agent)}

        # Advance current agent index
        self._current_agent_idx = self.agents.index(next_agent)

        return observations, rewards, terminations, truncations, infos

    def observe(self, agent: str) -> dict:
        """Get observation for a specific agent."""
        return self._get_obs(agent)

    def get_action_mask(self, agent: str) -> np.ndarray:
        """Get valid action mask for a specific agent."""
        player_idx = int(agent.split("_")[1])
        return self.game.get_action_mask(player_idx)

    def action_space(self, agent: str) -> spaces.Discrete:
        """Get action space for a specific agent."""
        return spaces.Discrete(self.game.num_actions)

    def observation_space(self, agent: str) -> spaces.Dict:
        """Get observation space for a specific agent."""
        return self.observation_space

    def render(self) -> str | None:
        """Render the current state."""
        if self.render_mode is None:
            return None

        from blokus_gym.utils.render import render_ansi_board, render_text_board

        if self.render_mode == "human":
            print(render_text_board(self.game))
            return None
        elif self.render_mode == "ansi":
            return render_ansi_board(self.game)
        else:
            raise ValueError(f"Unknown render_mode: {self.render_mode}")

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

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

    def _get_info(self, agent: str) -> dict:
        """Get auxiliary info for an agent."""
        player_idx = int(agent.split("_")[1])
        return {
            "action_mask": self.game.get_action_mask(player_idx),
            "current_player": self.game.current_player,
            "step_count": self._cur_step,
            "scores": self.game.get_scores(),
        }

    def agent_iter(self):
        """Iterate over agents in turn order."""
        # This is a generator that yields the current agent
        while True:
            agent = self.agents[self.game.current_player]
            if self._agent_dones.get(agent, False):
                # Skip done agents
                if all(self._agent_dones.values()):
                    break
                continue
            yield agent
