"""Vectorised CarRacing-v3 factory.

Wraps the existing single-env preprocessing stack inside a Gymnasium
vector env so n_envs copies can step in parallel processes (async mode).
"""

import gymnasium as gym

from src.env_wrappers import FrameStack, GrayScaleResize, SkipFrame


def _make_one(rank: int, seed: int):
  def _init():
      env = gym.make("CarRacing-v3", continuous=False)
      env = SkipFrame(env, k=4)
      env = GrayScaleResize(env, size=84)
      env = FrameStack(env, k=4)
      env.action_space.seed(seed + rank)
      return env
  return _init


def make_vec_env(n_envs: int = 4, seed: int = 0, async_mode: bool = True):
  """Build a vectorised CarRacing-v3 env.

  Returns a gym.vector env whose obs has shape (n_envs, 4, 84, 84) uint8.
  """
  fns = [_make_one(i, seed) for i in range(n_envs)]
  if async_mode:
      return gym.vector.AsyncVectorEnv(fns)
  return gym.vector.SyncVectorEnv(fns)