from __future__ import annotations

from typing import TYPE_CHECKING, cast

if TYPE_CHECKING:
    from pettingzoo.utils.env import ParallelEnv
    from stable_baselines3.common.vec_env import VecEnv

    from ..types import Action, Observation

import types


def make_vec_from_penv(env: ParallelEnv[int, Observation, Action],
                       num_envs: int = 1,
                       processes: int = 1,
                       seed: int = 1) -> VecEnv:
    """
    Creates an imitation-compatible vectorized enviroment from
    a :py:class:`PettingZoo  Parallel environment <pettingzoo.utils.env.ParallelEnv>`,
    similarly as :py:func:`navground.learning.parallel_env.make_vec_from_penv` but
    adding a (custom) ``RolloutInfoWrapper`` to collect rollouts.

    It first creates a :py:class:`supersuit.vector.MarkovVectorEnv`, and then
    applies :py:func:`supersuit.concat_vec_envs_v1` to concatenate ``number`` copies of it.

    :param      env:       The environment
    :param      num_envs:  The number of pettingzoo environments to stuck together
    :param      processes: The number of (parallel) processes
    :param      seed:      The seed

    :returns:   The vectorized environment.
    """
    import supersuit  # type: ignore[import-untyped]

    from .parallel_rollout_wrapper import RolloutInfoWrapper

    env = RolloutInfoWrapper(env)
    menv = supersuit.vector.MarkovVectorEnv(env, black_death=True)
    venv = supersuit.concat_vec_envs_v1(menv,
                                        num_envs,
                                        num_cpus=processes,
                                        base_class="stable_baselines3")

    def get_attr(self, attr_name: str, indices=None) -> list:
        if indices is None:
            indices = range(num_envs * menv.num_envs)
        if isinstance(indices, int):
            indices = [indices]
        rs = []
        for i in indices:
            penv = self.vec_envs[i // menv.num_envs].par_env
            rs.append(getattr(penv, attr_name))
        if attr_name == 'render_mode':
            # This is an attribute of vector environments themselves
            return rs[0]
        return rs

    def set_attr(self, attr_name: str, value, indices=None) -> None:
        if indices is None:
            indices = range(num_envs * menv.num_envs)
        if isinstance(indices, int):
            indices = [indices]
        for i in indices:
            penv = self.vec_envs[i // menv.num_envs].par_env
            setattr(penv, attr_name, value)

    cenv = venv.unwrapped
    cenv.get_attr = types.MethodType(get_attr, cenv)
    cenv.set_attr = types.MethodType(set_attr, cenv)

    for i, env in enumerate(venv.venv.vec_envs):
        env.reset(seed + i)
    return cast("VecEnv", venv)
