import gymnasium as gym
import gymnasium_robotics  # registers Mujoco robotics envs into gym
from envhub.envs.core.base_env import EnvHubEnv
from envhub.registry import register_env


def load_env(reward_type="dense", max_episode_steps=50, **kwargs):
    """
    Loader for MuJoCo FetchPickAndPlace environment.
    Args:
        reward_type: "dense" or "sparse"
        max_episode_steps: truncation length
        kwargs: passed to gym.make
    """

    # Register robotics envs into Gymnasium
    gym.register_envs(gymnasium_robotics)

    env_id = "HandManipulateEgg-v1" if reward_type == "dense" else "HandManipulateEgg-v1"
    env = gym.make(env_id, max_episode_steps=max_episode_steps, **kwargs)

    return EnvHubEnv(
        env=env,
        name="mujoco/hand_manipulate_egg",
        observation_space=env.observation_space,
        action_space=env.action_space,
    )