# Copyright (c) 2026, HUAWEI CORPORATION.  All rights reserved.

import glob
import logging
import os
import sys

import torch

_logger = logging.getLogger(__name__)

_SO_NAME = "mindspeed_ops_C"


def _load_mindspeed_ops_library():
    """Load the compiled mindspeed_ops_C shared library.

    The shared library is built by CMake and installed alongside this
    __init__.py. Loading it triggers TORCH_LIBRARY registration so that
    ``torch.ops.mindspeed_ops.*`` operators become available.

    pybind11_add_module produces a CPython ABI-tagged name like
    ``mindspeed_ops_C.cpython-311-aarch64-linux-gnu.so``, so we glob
    for any file starting with ``mindspeed_ops_C``.
    """
    so_dir = os.path.dirname(os.path.abspath(__file__))

    for pattern in (_SO_NAME + "*.so", _SO_NAME + "*.pyd"):
        candidates = glob.glob(os.path.join(so_dir, pattern))
        if candidates:
            so_path = candidates[0]
            torch.ops.load_library(so_path)
            _logger.info("Loaded mindspeed_ops library: %s", so_path)
            return

    # .so not found — the package was not built for this platform
    print(
        f"[mindspeed_ops] WARNING: {_SO_NAME}{{*.so,*.pyd}} not found in {so_dir}. "
        f"Rebuild on an NPU platform with: pip install -e .",
        file=sys.stderr,
    )


def _bootstrap_custom_op_env():
    """Set ASCEND_CUSTOM_OPP_PATH so CANN can discover custom AscendC kernels.

    Called automatically on ``import mindspeed_ops``, so users do not need to
    manually export environment variables.
    """
    so_dir = os.path.dirname(os.path.abspath(__file__))
    vendor_path = os.path.join(so_dir, "_cann_ops_custom", "vendors", "mindspeed_transformer")
    if not os.path.isdir(vendor_path):
        return

    current = os.environ.get("ASCEND_CUSTOM_OPP_PATH", "")
    entries = [e for e in current.split(":") if e]
    if vendor_path not in entries:
        entries.insert(0, vendor_path)
        os.environ["ASCEND_CUSTOM_OPP_PATH"] = ":".join(entries)


def _preload_custom_opapi_library():
    """Preload libcust_opapi.so into the process.

    torch_npu's aclnn wrapper resolves custom op entry points with
    ``dlopen("libcust_opapi.so")`` by bare name, which only searches the
    dynamic linker paths and NOT the package's ``_cann_ops_custom``
    directory, so the lookup is unreliable. Loading the library ourselves
    by absolute path ensures any later bare-name dlopen reuses this
    already-loaded object, making custom op resolution deterministic.

    Skipped silently when the library is absent (e.g. unbuilt source
    installs); load failures are logged but never break the import.
    """
    import ctypes

    so_dir = os.path.dirname(os.path.abspath(__file__))
    lib_path = os.path.join(
        so_dir,
        "_cann_ops_custom",
        "vendors",
        "mindspeed_transformer",
        "op_api",
        "lib",
        "libcust_opapi.so",
    )
    if not os.path.isfile(lib_path):
        return

    try:
        ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)
        _logger.info("Preloaded custom op api library: %s", lib_path)
    except OSError as exc:
        _logger.warning(
            "Failed to preload %s: %s. Custom aclnn ops may fail to resolve.",
            lib_path,
            exc,
        )


def _preload_catlass_libraries():
    """Preload bundled CATLASS DSOs so mindspeed_ops_C resolves them without RPATH.

    CATLASS-enabled wheels ship ``libmindspeed_catlass_*.so`` under the
    package's ``lib/`` directory, and the extension records them as NEEDED
    by soname. Loading them by absolute path first registers each object
    under its soname, so the dynamic linker reuses these already-loaded
    objects when ``mindspeed_ops_C`` is loaded afterwards. This replaces
    the former ``$ORIGIN/lib`` RPATH, which is forbidden by the secure
    compile rules.

    No-op when ``lib/`` does not exist (CATLASS-disabled builds).
    """
    import ctypes

    lib_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lib")
    if not os.path.isdir(lib_dir):
        return

    for so_path in sorted(glob.glob(os.path.join(lib_dir, "*.so"))):
        try:
            ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)
            _logger.info("Preloaded CATLASS library: %s", so_path)
        except OSError as exc:
            _logger.warning("Failed to preload %s: %s", so_path, exc)


_preload_catlass_libraries()
_load_mindspeed_ops_library()
_bootstrap_custom_op_env()
_preload_custom_opapi_library()