from __future__ import annotations

import numpy as np
from gymnasium import spaces
from gymnasium.core import Env, Wrapper


class ActionMaskWrapper(Wrapper):
    """Wrapper that moves the action mask from info into the observation space.

    This is useful for RL frameworks that expect the action mask to be part
    of the observation (e.g., Stable-Baselines3 with custom policies,
    RLlib, etc.).

    After wrapping, the observation becomes a Dict with:
        - "observation": the original observation
        - "action_mask": boolean array of valid actions

    Example:
        >>> from blokus_gym import BlokusEnv
        >>> from blokus_gym.wrappers import ActionMaskWrapper
        >>> env = BlokusEnv()
        >>> env = ActionMaskWrapper(env)
        >>> obs, info = env.reset()
        >>> obs["action_mask"]  # Boolean mask of valid actions
    """

    def __init__(self, env: Env):
        super().__init__(env)

        # Build new observation space
        original_obs_space = env.observation_space
        action_dim = env.action_space.n

        self.observation_space = spaces.Dict({
            "observation": original_obs_space,
            "action_mask": spaces.Box(0, 1, (action_dim,), dtype=bool),
        })

    def reset(self, *, seed=None, options=None):
        obs, info = self.env.reset(seed=seed, options=options)
        action_mask = info.get("action_mask", np.zeros(self.env.action_space.n, dtype=bool))
        return {"observation": obs, "action_mask": action_mask}, info

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        action_mask = info.get("action_mask", np.zeros(self.env.action_space.n, dtype=bool))
        return {"observation": obs, "action_mask": action_mask}, reward, terminated, truncated, info
