import gymnasium as gym
from envhub.core.base_env import EnvHubEnv
from envhub.registry import register_env

# MiniWoB imports
try:
    import miniwob
except ImportError:
    raise ImportError("Please install MiniWoB: `pip install miniwob`")

def load_env(task: str = "click-button", max_episode_steps: int = 50, **kwargs):
    """
    Loader for a MiniWoB browser interaction environment.

    Args:
        task (str): MiniWoB task name (e.g., "click-button", "book-flight", "choose-date").
        max_episode_steps (int): maximum number of steps per episode.
        kwargs: additional arguments passed to gym.make()

    Returns:
        EnvHubEnv: unified environment wrapper.
    """

    # compose the MiniWoB env ID (e.g. "miniwob/click-button-v1")
    env_id = f"miniwob/{task}-v1"

    # Make the environment
    env = gym.make(env_id, max_episode_steps=max_episode_steps, **kwargs)

    # Wrap it as an EnvHubEnv
    return EnvHubEnv(
        env=env,
        name=f"browser/{task}",
        observation_space=env.observation_space,
        action_space=env.action_space,
    )
register_env("browser/miniwob_click_button", load_env)