# Copyright (c) 2026 INSAIT. MIT License (see LICENSE).
"""LADDER SimplerEnv-Long for LeRobot EnvHub: 54 long-horizon WidowX tasks on SAPIEN 3 / ManiSkill 3.

Each task is one suite of LeRobot's ``{suite: {task_id: env}}``; its 24 episodes are the 24 object
layouts LADDER defines (24 parallel envs on one GPU). See README.md.
"""

import os
import re
import sys

REPO_ID = "INSAIT-Institute/LADDER-SimplerEnv-Long"

_HERE = os.path.dirname(os.path.abspath(__file__))
if not os.path.exists(os.path.join(_HERE, "_ladder_envhub.py")):
    # `make_env("<repo>")` downloads only this file: fetch the rest of the same revision.
    from huggingface_hub import snapshot_download

    _revision = os.path.basename(_HERE)
    _HERE = snapshot_download(REPO_ID, revision=_revision if re.fullmatch(r"[0-9a-f]{40}", _revision) else None)
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

from functools import partial  # noqa: E402
from typing import Any  # noqa: E402

import _ladder_envhub as lh  # noqa: E402

EPISODES_PER_TASK = 24
# ManiSkill reads its asset directory once, when first imported; the tree is filled on the first reset.
_MANISKILL_DIR = lh.cache_dir() / "maniskill3"
os.environ.setdefault("MS_ASSET_DIR", str(_MANISKILL_DIR))
LEROBOT_SIM = lh.lerobot_sim("widowx", fps=5)


def task_names() -> list[str]:
    """The suite's task names, from LADDER's own registry."""
    from ladder.components.inference.robots.simpler.composition.specs import SPECS

    return sorted(spec.name for spec in SPECS)


class SimplerEnvLong(lh.LadderVectorEnv):
    suite = "simpler"
    robot = "widowx"
    resizable = False  # LADDER's SimplerEnv server renders its 480x640 camera
    episodes_per_task = EPISODES_PER_TASK
    fps = 5

    @classmethod
    def task_budget(cls, task: str) -> int:
        from ladder.components.inference.robots.simpler.composition.specs import SPECS

        return next(spec.max_steps for spec in SPECS if spec.name == task)

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._env = None
        # Once per process, before the server reports ready; the scene is built on the first reset.
        if os.environ["MS_ASSET_DIR"] == str(_MANISKILL_DIR):
            _stage_maniskill_assets()
        _seed_sapien_physx()
        _point_vulkan_at_nvidia()

    def _backend_reset(self, episode_ids: list[int]) -> list[Any]:
        if self._env is None:
            from ladder.components.inference.robots.simpler.env_builder import ABS_CONTROL_MODE
            from ladder.components.inference.robots.simpler.simpler_env import SimplerEnv, SimplerEnvConfig

            # The settings of LADDER's server config, robots/simpler/configs/widowx_composition.yaml.
            self._env = SimplerEnv(
                SimplerEnvConfig(
                    env_name=f"{self.task}-v0",
                    robot_name="widowx",
                    scene_name="bridge_table_1_v1",
                    rgb_overlay_path="",  # ManiSkill3 picks the overlay from env_name
                    sim_freq=500,
                    control_mode=ABS_CONTROL_MODE,
                    control_freq=5,
                    max_episode_steps=0,  # the task's own budget, 100 steps per subgoal
                    raytracing=False,
                    n_envs=self.num_envs,
                    episode_layouts=EPISODES_PER_TASK,
                )
            )
        return self._env.reset([e % EPISODES_PER_TASK for e in episode_ids])

    def _backend_step(self, actions: list[Any]) -> list[Any]:
        return self._env.step(actions)

    def _backend_close(self) -> None:
        if self._env is not None:
            self._env.close()
            self._env = None


def _stage_maniskill_assets() -> None:
    """ManiSkill reads one asset tree: LADDER's bridge scenes and ManiSkill's WidowX model side by side."""
    merged = _MANISKILL_DIR
    if not (merged / ".complete").exists():
        assets, widowx = lh.fetch("simpler_assets"), lh.fetch("widowx")
        (merged / "data" / "robots").mkdir(parents=True, exist_ok=True)
        for entry in (assets / "data").iterdir():
            link = merged / "data" / entry.name
            if entry.name != "robots":
                _link(link, entry)
        for entry in (assets / "data" / "robots").iterdir() if (assets / "data" / "robots").is_dir() else ():
            _link(merged / "data" / "robots" / entry.name, entry)
        _link(merged / "data" / "robots" / "widowx", widowx / "data" / "robots" / "widowx")
        (merged / ".complete").touch()


def _link(link, target) -> None:
    if not link.is_symlink():
        try:
            link.symlink_to(target)
        except FileExistsError:  # another process made it
            pass


def _seed_sapien_physx() -> None:
    """SAPIEN downloads its PhysX GPU library from GitHub on first use; use LADDER's pinned copy."""
    import sapien.physx

    library = os.path.join(os.path.expanduser("~"), ".sapien", "physx", sapien.physx.version(), "libPhysXGpu_64.so")
    if not os.path.exists(library):
        os.makedirs(os.path.dirname(library), exist_ok=True)
        temporary = f"{library}.{os.getpid()}"
        os.symlink(lh.fetch("sapien_physx") / "libPhysXGpu_64.so", temporary)
        os.replace(temporary, library)


def _point_vulkan_at_nvidia() -> None:
    """SAPIEN renders with Vulkan. A container may have the NVIDIA driver but no ICD file naming it."""
    if os.environ.get("VK_ICD_FILENAMES"):
        return
    for directory in ("/etc/vulkan/icd.d", "/usr/share/vulkan/icd.d"):
        if os.path.exists(os.path.join(directory, "nvidia_icd.json")):
            return
    icd = lh.cache_dir() / "nvidia_icd.json"
    if not icd.exists():
        icd.write_text('{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3.0"}}')
    os.environ["VK_ICD_FILENAMES"] = str(icd)


def vector_env_fns(
    task: str = "",
    n_envs: int = EPISODES_PER_TASK,
    **options: Any,
) -> dict[str, dict[int, Any]]:
    """One factory per task, ``{task: {0: factory}}``: every task, or the comma-separated ``task``.

    LeRobot builds only the tasks it runs (``--env.task``), so listing all 54 costs nothing.
    """
    names = lh.select_tasks(task, task_names())
    if os.environ["MS_ASSET_DIR"] == str(_MANISKILL_DIR):
        _stage_maniskill_assets()  # download before the server reports ready
    lh.fetch("sapien_physx")
    return {name: {0: partial(SimplerEnvLong, name, n_envs, **options)} for name in names}


def make_env(n_envs: int = EPISODES_PER_TASK, use_async_envs: bool = False, task: str = "", **options: Any):
    """One task as a batched vector env (``--env.use_native_batch``); LeRobot otherwise uses ``vector_env_fns``."""
    name = lh.one_task(task, task_names())
    return {name: {0: vector_env_fns(name, n_envs, **options)[name][0]()}}