"""Resolve release assets relative to the repository root.

The H2 release includes its MJCF/meshes, native CLoSD bundle, ready pose,
Stage-C policy and residual ONNX. Legacy helpers can use ICRA_REPO_ROOT
when explicitly working with a separate training checkout; those helpers
are not required by the released SHM MuJoCo runner.
"""

from __future__ import annotations

import importlib.util
import os
import sys
import types
from functools import lru_cache
from pathlib import Path

# These core files live under ``deploy_closd/runtime`` in the real-robot
# bundle, while all models and robot assets live one level above them.
PACKAGE_DIR = Path(__file__).resolve().parents[1]


def package_dir() -> Path:
    return PACKAGE_DIR


def resolve(relative: str | Path) -> Path:
    """Resolve a bundle-relative path; fall back to the training checkout when absent."""
    path = Path(relative)
    if path.is_absolute():
        return path
    local = PACKAGE_DIR / path
    if local.exists():
        return local
    repo = _repo_root_optional()
    if repo is not None:
        for candidate in (repo / path, repo / "source" / "whole_body_tracking" / "whole_body_tracking" / path):
            if candidate.exists():
                return candidate
    return local


def _repo_root_optional() -> Path | None:
    override = os.environ.get("ICRA_REPO_ROOT", "").strip()
    candidates = [Path(override).expanduser()] if override else []
    candidates += [PACKAGE_DIR.parent / "humanoid_table_tennis"]
    for cand in candidates:
        if (cand / "source" / "whole_body_tracking").is_dir():
            return cand.resolve()
    return None


def assets_dir() -> Path:
    return PACKAGE_DIR / "assets" / "robot_description"


def sim2sim_dir() -> Path:
    local = PACKAGE_DIR / "sim2sim"
    if local.is_dir():
        return local
    repo = _repo_root_optional()
    if repo is None:
        raise FileNotFoundError("deploy/sim2sim helper modules are missing and no training checkout is available")
    return repo / "sim2sim_tt"


def ensure_sim2sim_on_path() -> None:
    path = str(sim2sim_dir())
    if path not in sys.path:
        sys.path.insert(0, path)


def runtime_dir() -> Path:
    local = PACKAGE_DIR / "runtime"
    if (local / "motionbricks_closd_runtime.py").is_file():
        return local
    repo = _repo_root_optional()
    if repo is None:
        raise FileNotFoundError("deploy/runtime modules are missing and no training checkout is available")
    return repo / "source" / "whole_body_tracking" / "whole_body_tracking" / "tasks" / "tracking" / "mdp"


@lru_cache(maxsize=1)
def load_closd_runtime_module() -> types.ModuleType:
    """Load ``motionbricks_closd_runtime`` (native-bundle path) without importing the training package."""
    mdp_dir = runtime_dir()
    pkg_name = "_wbt_closd_runtime"
    if pkg_name in sys.modules:
        return sys.modules[f"{pkg_name}.motionbricks_closd_runtime"]
    pkg = types.ModuleType(pkg_name)
    pkg.__path__ = [str(mdp_dir)]  # relative imports inside the files resolve against this
    sys.modules[pkg_name] = pkg
    for name in ("native_closd_inference", "motionbricks_closd_runtime"):
        spec = importlib.util.spec_from_file_location(f"{pkg_name}.{name}", mdp_dir / f"{name}.py")
        module = importlib.util.module_from_spec(spec)
        sys.modules[f"{pkg_name}.{name}"] = module
        spec.loader.exec_module(module)  # type: ignore[union-attr]
    return sys.modules[f"{pkg_name}.motionbricks_closd_runtime"]