"""Shared loader helper for the examples.

Provides `load_kernel()` which returns the mamba3 kernel package. Uses the
local clone (examples/../build/torch-neuron/) when running from a git clone,
or falls back to the HF Hub via `get_kernel("jburtoft/mamba3-neuron-kernels")`
when installed via the `kernels` library.
"""

import os
import sys


def load_kernel():
    """Return the mamba3 kernel module."""
    local_path = os.path.abspath(
        os.path.join(os.path.dirname(__file__), "..", "build", "torch-neuron", "__init__.py")
    )
    if os.path.exists(local_path):
        # print(f"[loader] using local clone: {os.path.dirname(local_path)}")
        import importlib.util
        spec = importlib.util.spec_from_file_location(
            "mamba3", local_path,
            submodule_search_locations=[os.path.dirname(local_path)],
        )
        m = importlib.util.module_from_spec(spec)
        sys.modules["mamba3"] = m
        spec.loader.exec_module(m)
        return m
    else:
        # print("[loader] using HF Hub: jburtoft/mamba3-neuron-kernels")
        from kernels import get_kernel
        return get_kernel(
            "jburtoft/mamba3-neuron-kernels",
            revision="v1.0.0",
            trust_remote_code=True,
        )