# Copyright (c) 2026 INSAIT. MIT License (see LICENSE).
"""LADDER RoboLab-Long for LeRobot EnvHub: 44 long-horizon Franka (DROID) tasks on Isaac Sim 6 / Isaac Lab 3.

Each task is one suite of LeRobot's ``{suite: {task_id: env}}``; its 8 episodes run as 8 parallel envs
in one Isaac Lab world. Isaac Sim starts once per process. See README.md.
"""

import os
import re
import sys

REPO_ID = "INSAIT-Institute/LADDER-RoboLab-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)

import atexit  # noqa: E402
import logging  # noqa: E402
from concurrent.futures import ThreadPoolExecutor  # noqa: E402
from functools import partial  # noqa: E402
from typing import Any  # noqa: E402

import _ladder_envhub as lh  # noqa: E402

EPISODES_PER_TASK = 8
CONTROL_HZ = 15
LEROBOT_SIM = lh.lerobot_sim("franka", fps=CONTROL_HZ)


def _add_sources() -> None:
    """Isaac Lab, RoboLab (patched to read its assets from ROBOLAB_ASSET_DIR) and openpi_client, as LADDER pins them."""
    if _Isaac.sources:
        return
    robolab = lh.fetch("robolab_src")
    constants = robolab / "robolab" / "constants.py"
    text = constants.read_text()
    patched = 'ASSET_DIR = os.environ.get("ROBOLAB_ASSET_DIR") or os.path.join(PACKAGE_DIR, "assets")'
    if patched not in text:
        # LADDER's build_support/python/patches/robolab_constants_asset_dir.patch.
        text = text.replace('ASSET_DIR = os.path.join(PACKAGE_DIR, "assets")', patched, 1)
        text = text.replace(
            "    return os.path.join(PACKAGE_DIR, relative_path)",
            '    normalized = relative_path.replace(os.sep, "/")\n'
            '    if normalized.startswith("assets/"):\n'
            '        return os.path.join(ASSET_DIR, *normalized[len("assets/"):].split("/"))\n'
            "    return os.path.join(PACKAGE_DIR, relative_path)",
            1,
        )
        # Atomically: another process may be importing it.
        staged = constants.with_name(f".constants.{os.getpid()}.py")
        staged.write_text(text)
        os.replace(staged, constants)
    if patched not in constants.read_text():
        raise RuntimeError(f"{constants}: LADDER's asset-directory patch no longer applies")
    # The extracted wheel is an installed package on sys.path, with its dist-info (see _ladder_envhub.DOWNLOADS).
    lh.add_to_path(lh.fetch("isaaclab"), robolab, lh.fetch("openpi_src") / "packages" / "openpi-client" / "src")
    os.environ.setdefault("ROBOLAB_ASSET_DIR", str(lh.fetch("robolab_assets") / "assets"))
    _Isaac.sources = True


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

    return sorted(SPECS)


class _Isaac:
    """Isaac Sim starts once per process and runs on one thread; one task's world is kept at a time."""

    thread = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ladder-isaac")
    sources = False
    app: Any = None
    env: Any = None
    env_cfg: Any = None
    oracle: Any = None
    key: Any = None


def _start_isaac() -> None:
    if _Isaac.app is not None:
        return
    _add_sources()
    # Running this environment accepts the NVIDIA Isaac Sim EULA (see README); Kit asks on stdin otherwise.
    os.environ.setdefault("OMNI_KIT_ACCEPT_EULA", "YES")
    os.environ.setdefault("ACCEPT_EULA", "Y")
    import cv2  # noqa: F401
    from isaaclab.app import AppLauncher

    root = logging.getLogger()
    levels = {handler: handler.level for handler in root.handlers}
    with lh.ignore_signal_in_thread():
        _Isaac.app = AppLauncher(headless=True, enable_cameras=True).app
    for handler, level in levels.items():  # Kit lowers the levels of existing handlers
        handler.setLevel(level)
    # Kit hangs in its own shutdown; leave without it, keeping the exit status.
    atexit.register(
        lambda: (
            sys.stdout.flush(),
            sys.stderr.flush(),
            os._exit(0 if getattr(sys, "last_value", None) is None else 1),
        )
    )
    import carb.settings
    import robolab.constants

    # As LADDER's env_server.py: LADDER scores subgoals itself, and renderer failures stay visible.
    robolab.constants.ENABLE_SUBTASK_PROGRESS_CHECKING = False
    carb.settings.get_settings().set("/persistent/app/usd/muteUsdDiagnostics", False)


class RoboLabLong(lh.LadderVectorEnv):
    suite = "robolab"
    robot = "franka"
    default_image_size = (224, 224)
    episodes_per_task = EPISODES_PER_TASK
    fps = CONTROL_HZ

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

        return int(SPECS[task].episode_length_s * CONTROL_HZ)

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._env = None
        self._raw = None
        self._converter = None
        # LeRobot's server builds an env (one env of the first task) before it reports ready. Isaac Sim starts here,
        # and the first env of a process builds its world here too: that compiles the shaders and loads the assets
        # every task shares, which on a new machine or container would otherwise fall in a client's first open.
        # A client's own batch size gets its world on its first reset.
        self._fresh = _Isaac.thread.submit(self._start).result()

    def _start(self) -> bool:
        _start_isaac()
        return _Isaac.env is None and self._build()

    def goal_sources(self) -> list[str | None] | None:
        # The object each goal places (a relay's legs are separate goals), as LADDER's privileged.py names it.
        return [
            next((geometry[0] for _, geometry in reversed(list(goal.legs)) if geometry is not None), None)
            for goal in _Isaac.oracle or []
        ]

    def _backend_reset(self, episode_ids: list[int]) -> list[Any]:
        return _Isaac.thread.submit(self._reset).result()

    def _backend_step(self, actions: list[Any]) -> list[Any]:
        if _Isaac.env is not self._env:
            raise RuntimeError("Another task's world replaced this one: run one task at a time (--env.max_parallel_tasks=1)")
        return _Isaac.thread.submit(self._step, actions).result()

    def _build(self) -> bool:
        """Build this task's world at this batch size, unless the process holds it; True if it was built."""
        from ladder.components.inference.robots.robolab.composition.specs import task_class_name, task_target
        from ladder.components.inference.robots.robolab.oracle import build_oracle
        from robolab.core.environments.runtime import create_env
        from robolab.registrations.droid.auto_env_registrations_jointpos import auto_register_droid_envs

        self._class_name = task_class_name(self.task)
        key = (self._class_name, self.num_envs)
        if _Isaac.key == key:
            return False
        if _Isaac.env is not None:
            _Isaac.env.close()
            _Isaac.env = _Isaac.key = None
        auto_register_droid_envs(task_dirs=None, task=[task_target(self.task)], cameras=None)
        # The defaults of LADDER's env_server.py (RoboLab's add_common_eval_args).
        _Isaac.env, _Isaac.env_cfg = create_env(
            self._class_name,
            device="cuda:0",
            num_envs=self.num_envs,
            instruction_type="default",
            policy="ladder_runner",
            renderer="realtime",
            rendering_mode="balanced",
        )
        _Isaac.oracle = build_oracle(getattr(_Isaac.env.cfg, "subtasks", None), _Isaac.env_cfg.instruction)
        _Isaac.key = key
        if int(_Isaac.env.max_episode_length) != self._max_episode_steps:
            logging.getLogger(__name__).warning(
                "%s: Isaac Lab budget %d steps, expected %d", self.task, _Isaac.env.max_episode_length,
                self._max_episode_steps,
            )
        return True

    def _reset(self) -> list[Any]:
        from ladder.components.inference.robots.robolab.droid_conversions import DroidJointposConverter

        # As LADDER's env_server.py: a new world is reset as it is, a kept one first unfreezes its finished envs.
        if not self._build() and not self._fresh:
            _Isaac.env.reset_eval_state()
        self._fresh = False
        self._env = _Isaac.env
        self._env.reset()
        self._raw, _ = self._env.reset()  # RoboLab's own double reset (robolab.eval.episode.run_episode)
        # DroidJointposConverter takes (width, height), as LADDER's --image-size.
        self._converter = DroidJointposConverter(num_envs=self.num_envs, image_size=(self._w, self._h))
        return self._observations()

    def _step(self, actions: list[Any]) -> list[Any]:
        import omni.kit.app
        import omni.timeline
        import torch

        joints = self._raw["proprio_obs"]["arm_joint_pos"].detach().cpu().numpy()
        commands = self._converter.actions_to_jointpos(actions, joints)
        timeline, kit = omni.timeline.get_timeline_interface(), omni.kit.app.get_app()
        while not timeline.is_playing():
            kit.update()
        self._raw, _, _, _, _ = self._env.step(torch.as_tensor(commands, device=self._env.device))
        return self._observations()

    def _observations(self) -> list[Any]:
        images, proprio = self._raw["image_obs"], self._raw["proprio_obs"]
        results = self._env.get_env_results()
        active = set(self._env.active_env_ids)
        # During a step, _step_count is not yet incremented; LADDER's server counts the step first.
        out_of_time = self._step_count + 1 >= int(self._env.max_episode_length)
        achieved = None
        if _Isaac.oracle:
            achieved = [[bool(g.check(self._env, env_id=i)) for g in _Isaac.oracle] for i in range(self.num_envs)]
        observations = []
        for i in range(self.num_envs):
            success = bool(results[i]["success"])
            metadata = {
                "env_name": self._class_name,
                "robot_name": "franka",
                "instruction": _Isaac.env_cfg.instruction,
                "done": success,
                "truncated": out_of_time or i not in active,
                "oracle": {"achieved": achieved[i] if achieved is not None else [success]},
            }
            observations.append(
                self._converter.observation_dict(
                    env_i=i,
                    main_image=images["over_shoulder_left_camera"][i].detach().cpu().numpy(),
                    wrist_image=images["wrist_cam"][i].detach().cpu().numpy(),
                    joints=proprio["arm_joint_pos"][i].detach().cpu().numpy(),
                    gripper_position=proprio["gripper_pos"][i].detach().cpu().numpy(),
                    metadata=metadata,
                )
            )
        return observations


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 44 costs nothing.
    """
    names = lh.select_tasks(task, task_names())
    _add_sources()  # download before the server reports ready
    return {name: {0: partial(RoboLabLong, 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]()}}