#!/usr/bin/env python3
"""启动 Isaac Sim 后验证全部 Gen2 Task 注册和 Train/Play 契约。"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = REPO_ROOT / "source"
if str(SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(SOURCE_ROOT))

from isaaclab.app import AppLauncher


# Isaac Lab 配置模块依赖 pxr，必须先启动 App，再导入 engineai_lab.tasks。
app_launcher = AppLauncher(headless=True)
simulation_app = app_launcher.app


def _load_class(entry_point: str):
    module_name, class_name = entry_point.rsplit(":", 1)
    return getattr(importlib.import_module(module_name), class_name)


def main() -> None:
    import gymnasium as gym

    import engineai_lab.tasks  # noqa: F401
    from engineai_lab.tasks.velocity.config.gen2.agents.rsl_rl_ppo_cfg import (
        Gen2FastPPORunnerCfg as LegacyFastRunnerCfg,
    )
    from engineai_lab.tasks.velocity.config.gen2.agents.fast_ppo_cfg import Gen2FastPPORunnerCfg
    from engineai_lab.tasks.velocity.config.gen2.flat_env_cfg import Gen2FastEnvCfg as LegacyFastEnvCfg
    from engineai_lab.tasks.velocity.config.gen2.registry import GEN2_TASK_SPECS
    from engineai_lab.tasks.velocity.config.gen2.stages.fast import Gen2FastEnvCfg

    # 旧模块现在是兼容层；类对象必须与新阶段模块完全相同。
    assert LegacyFastEnvCfg is Gen2FastEnvCfg
    assert LegacyFastRunnerCfg is Gen2FastPPORunnerCfg

    configs = {}
    for task_id, _, _ in GEN2_TASK_SPECS:
        spec = gym.spec(task_id)
        env_cfg_cls = _load_class(spec.kwargs["env_cfg_entry_point"])
        runner_cfg_cls = _load_class(spec.kwargs["rsl_rl_cfg_entry_point"])
        env_cfg = env_cfg_cls()
        runner_cfg = runner_cfg_cls()
        configs[task_id] = env_cfg

        assert runner_cfg.experiment_name == "velocity_flat_terrain_gen2"
        assert len(env_cfg.actions.joint_pos.joint_names) == 28
        assert env_cfg.observations.policy.joint_pos.history_length == 15
        print(f"OK {task_id}: {env_cfg_cls.__module__}.{env_cfg_cls.__name__}", flush=True)

    train_play_pairs = (
        ("Flat-Gen2-v0", "Flat-Gen2-Play-v0"),
        ("Flat-Gen2-Speed-v0", "Flat-Gen2-Speed-Play-v0"),
        ("Flat-Gen2-Natural-v0", "Flat-Gen2-Natural-Play-v0"),
        ("Flat-Gen2-Fast-v0", "Flat-Gen2-Fast-Play-v0"),
        ("Flat-Gen2-Sprint-v0", "Flat-Gen2-Sprint-Play-v0"),
        ("Flat-Gen2-NaturalRun-v0", "Flat-Gen2-NaturalRun-Play-v0"),
    )
    for train_id, play_id in train_play_pairs:
        train_cfg = configs[train_id]
        play_cfg = configs[play_id]
        assert train_cfg.actions.joint_pos.joint_names == play_cfg.actions.joint_pos.joint_names
        assert train_cfg.actions.joint_pos.scale == play_cfg.actions.joint_pos.scale

    print(f"RESULT=PASS tasks={len(configs)} train_play_pairs={len(train_play_pairs)}", flush=True)


if __name__ == "__main__":
    main()