"""
Environment factory for creating Gymnasium environments.

Handles:
- Environment instantiation
- Seeding for reproducibility
- Episode statistics tracking
- Vectorization for parallel environments
"""

import gymnasium as gym
from typing import Callable


def make_env(env_id: str, seed: int, idx: int = 0) -> Callable:
    """
    Create a single environment instance.
    
    Wraps the environment with:
    - RecordEpisodeStatistics: Track returns and lengths
    - Proper seeding for reproducibility
    
    Args:
        env_id: Gymnasium environment ID
        seed: Base random seed
        idx: Environment index (for parallel envs)
        
    Returns:
        Function that creates the environment
    """
    def thunk():
        env = gym.make(env_id)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env.action_space.seed(seed + idx)
        env.observation_space.seed(seed + idx)
        return env
    return thunk


def make_vec_envs(env_id: str, num_envs: int, seed: int) -> gym.vector.VectorEnv:
    """
    Create vectorized environments for parallel rollouts.
    
    Args:
        env_id: Gymnasium environment ID
        num_envs: Number of parallel environments
        seed: Base random seed
        
    Returns:
        Vectorized environment
    """
    envs = gym.vector.SyncVectorEnv(
        [make_env(env_id, seed, i) for i in range(num_envs)]
    )
    return envs