import asyncio

from browsergym.core.registration import register_task
from gymnasium import register

from webarenasafe.browser_env.custom_env import BrowserEnv
# register the WebArena benchmark
from . import config, task

ALL_WEBARENA_TASK_IDS = []


class WebArenaSafeEnv(BrowserEnv):  # Inherit from the existing environment
    def __init__(self, *args, **kwargs):
        print("custom WebArenaSafeEnv initialized has been called")

        super().__init__(*args, **kwargs,)
        # Your initialization here (e.g., set up Playwright)


# register the WebArena benchmark
for task_id in config.TASK_IDS:
    gym_id = f"browsergym/WebArenaSafeEnv.{task_id}"
    register(
        id=gym_id,  # Unique ID for your environment
        entry_point=lambda *env_args, **env_kwargs: WebArenaSafeEnv(task.GenericWebArenaTask, *env_args, **env_kwargs),
        # Replace with actual path
        nondeterministic=True,
        kwargs={"task_kwargs": {"task_id": task_id}},
    )
    # register_task()
    ALL_WEBARENA_TASK_IDS.append(gym_id)
