from typing import Literal, Any, Optional

import gymnasium as gym


def make_env_factory(
    gym_id: str,
    gym_type: Literal["gymnasium", "gym-v26", "gym-v21"] = "gymnasium",
    gym_vars: dict[str, Any] = {},
):
    def env_factory(index=0, **kwargs):
        env_vars = {
            **(gym_vars or {}),
            "render_mode": "rgb_array",
            # "index": index,
            **kwargs,
        }

        env = None
        if gym_type == "gymnasium":
            env = gym.make(gym_id, **env_vars)
        elif gym_type == "gym-v26":
            env = gym.make("GymV26Environment-v0", env_id=gym_id, **env_vars)
        elif gym_type == "gym-v21":
            env = gym.make("GymV21Environment-v0", env_id=gym_id, **env_vars)
        else:
            raise EnvironmentError(
                f"Unsupported environment type: {gym_type}. "
                "Please use a supported environment."
            )

        return env

    return env_factory


def make_env(
    gym_id: str,
    gym_type: Literal["gymnasium", "gym-v26", "gym-v21"] = "gymnasium",
    *,
    render_mode: Optional[str] = None,
    **kwargs,
):
    return make_env_factory(gym_id, gym_type)(render_mode=render_mode, **kwargs)
