"""Development loader: JIT-builds the kernel and exposes the packaged API.

Binds torch.ops to the package's `_ops` slot and imports torch-ext/mamba3
unchanged, so local runs exercise the code that ships.
"""

import importlib.util
import sys
import types
from pathlib import Path

import torch
from torch.utils.cpp_extension import load

_ROOT = Path(__file__).parent
_NAME = "mamba3_cuda"

load(
    name=_NAME,
    sources=[str(_ROOT / "torch-ext" / "torch_binding.cpp"),
             str(_ROOT / "mamba3_cuda" / "mamba3_step.cu"),
             str(_ROOT / "mamba3_cuda" / "mamba3_fwd.cu"),
             str(_ROOT / "mamba3_cuda" / "mamba3_bwd.cu")],
    extra_include_paths=[str(_ROOT / "torch-ext")],
    # build.toml declares no CUDA flags, so the published variants compile with
    # nvcc defaults; keep the dev build on the same arithmetic.
    extra_cuda_cflags=["-O3"],
    is_python_module=False,
    verbose=False,
)

_ops_mod = types.ModuleType("_m3pkg._ops")
_ops_mod.ops = getattr(torch.ops, _NAME)
sys.modules["_m3pkg._ops"] = _ops_mod

_spec = importlib.util.spec_from_file_location(
    "_m3pkg", _ROOT / "torch-ext" / "mamba3" / "__init__.py",
    submodule_search_locations=[str(_ROOT / "torch-ext" / "mamba3")])
_pkg = importlib.util.module_from_spec(_spec)
sys.modules["_m3pkg"] = _pkg
_spec.loader.exec_module(_pkg)

def ext_ops():
    return getattr(torch.ops, _NAME)


forward = _pkg.forward
DecodeState = _pkg.DecodeState
cumulative_angles = _pkg.cumulative_angles
chunk_decay = _pkg.chunk_decay
dispatch_paths = _pkg.dispatch_paths
__all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay",
           "dispatch_paths"]