# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import enum
import logging
import os
import ctypes
import re
from pathlib import Path
from typing import Optional, Union

import nvshmem.core
import nvshmem.bindings as bindings
from nvshmem.core.nvshmem_types import *
import nvshmem.core.utils as utils
import nvshmem.core.memory as memory
from nvshmem import __version__
from nvshmem.core._internal_tracking import _mr_references, _cached_device, _debug_mode, InternalInitStatus

from cuda.pathfinder import load_nvidia_dynamic_lib, find_nvidia_header_directory
from cuda.core import Buffer, MemoryResource
from cuda.core import Device, system
from cuda.core import ObjectCode

import numpy as np

Comm = None
mpi = None
_mpi4py_enabled = False


def _import_mpi():
    global mpi, Comm, _mpi4py_enabled
    if _mpi4py_enabled:
        return
    try:
        import mpi4py.MPI as mpi_module
        from mpi4py.MPI import Comm
        mpi = mpi_module
        Comm = Comm
        _mpi4py_enabled = True
        import sys
        print(f"MPI4Py enabled: {_mpi4py_enabled}", file=sys.stderr, flush=True)
        logger.info(f"MPI4Py enabled: {_mpi4py_enabled}")
    except (ImportError, RuntimeError) as e:
        import sys
        print(f"MPI4Py not enabled: {e}", file=sys.stderr, flush=True)
        logger.warning(f"MPI4Py not enabled: {e}")
        _mpi4py_enabled = False


__all__ = [
    'get_unique_id', 'init', 'finalize', 'get_version', 'module_init', 'module_finalize', 'library_init',
    'library_finalize', 'UniqueID', 'DeviceLibLanguage', 'find_device_bitcode_library'
]


class DeviceLibLanguage(enum.Enum):
    """
    Language/format of the NVSHMEM device library.

    ``LTOIR`` maps to the installed ``libnvshmem_device.ltoir.fatbin`` entry point.
    On CUDA 13+ that file is a multi-arch fatbin created with
    ``fatbinary --image3=kind=nvvm``. On CUDA 12.8-12.x the install layout keeps
    the same ``.ltoir.fatbin`` suffix for consistency, but the file is a renamed
    single-arch ``sm_90`` LTOIR image rather than a true multi-arch fatbin.
    Consumers that require per-arch LTOIR on CUDA 12 should use the installed
    ``libnvshmem_device_sm_<arch>.ltoir`` file directly.
    """
    LLVM_BITCODE = "bc"
    LTOIR = "ltoir"
    STATIC = "a"


logger = logging.getLogger("nvshmem")

UniqueID = bindings.uniqueid


def get_version() -> Version:
    """
    Get the NVSHMEM4Py version
   
    Returns an object of type nvshmem.core.Version which is a Python class
    This class contains several strings which represent versions related to NVSHMEM

    ``Version.openshmem_spec_version`` is the OpenSHMEM Spec that this NVSHMEM was built against

    ``Version.nvshmem4py_version`` is the version of the NVSHMEM4Py python library

    ``Version.libnvshmem_version`` is the version of NVSHMEM library that this package has opened
    """
    # Load the nvshmemem host library
    # Must happen before we use bindings
    # We use this explicilty here because user is allowed to call get_version before init is called
    # Repeat calls during initwill be no-ops
    load_nvidia_dynamic_lib("nvshmem_host")
    # Spec version
    spec_major = ctypes.c_int()
    spec_minor = ctypes.c_int()
    bindings.info_get_version(ctypes.addressof(spec_major), ctypes.addressof(spec_minor))

    # libnvshmem version
    lib_major = ctypes.c_int()
    lib_minor = ctypes.c_int()
    lib_patch = ctypes.c_int()
    bindings.vendor_get_version_info(ctypes.addressof(lib_major), ctypes.addressof(lib_minor),
                                     ctypes.addressof(lib_patch))

    return Version(openshmem_spec_version=f"{int(spec_major.value)}.{int(spec_minor.value)}",
                   nvshmem4py_version=__version__,
                   libnvshmem_version=f"{int(lib_major.value)}.{int(lib_minor.value)}.{lib_patch.value}")


def get_unique_id(empty=False) -> UniqueID:
    """
    Retrieve or create a unique ID used for UID-based NVSHMEM initialization.

    This function wraps the underlying NVSHMEM binding for obtaining a unique ID
    required when using the ``uid`` initializer method. Only a single rank (typically rank 0)
    should call this with `empty=False` to generate the ID. Other ranks should call it
    with ``empty=True`` and receive the ID through a user-defined communication mechanism
    (e.g., MPI broadcast or socket transfer).

    Args:
        empty (bool): If True, returns an empty (uninitialized) unique ID structure.
            If False, calls the underlying NVSHMEM function to generate a valid unique ID.

    Returns:
        ``UniqueID``: A `UniqueID` object containing the generated or empty NVSHMEM unique ID.

    Raises:
       ``NvshmemError``: If retrieving the unique ID from NVSHMEM fails.

    Example:
        >>> if rank == 0:
        ...     uid = nvshmem.core.get_unique_id()
        ... else:
        ...     uid = nvshmem.core.get_unique_id(empty=True)
        ...
        >>> nvshmem.core.init(uid=uid, rank=rank, nranks=size, initializer_method="uid")
    """
    # Load the nvshmemem host library
    # Must happen before we use bindings
    # We use this explicilty here because user is allowed to call get_unique_id before init is called
    # Repeat calls during initwill be no-ops
    load_nvidia_dynamic_lib("nvshmem_host")

    unique_id = bindings.uniqueid()
    if empty:
        return unique_id
    # This function is in TRANSPARENT mode
    # The binding will check status
    bindings.get_uniqueid(unique_id.ptr)
    return unique_id


def init(device: Device = None,
         uid: bindings.uniqueid = None,
         rank: int = None,
         nranks: int = None,
         mpi_comm: Comm = None,
         initializer_method: str = "") -> None:
    """
    Initialize the NVSHMEM runtime with either MPI or UID-based bootstrapping.

    Args:
        - device (``cuda.core.Device``, required): A Device() that will be bound to this process. All NVSHMEM operations on this process will use this Device
        - uid (``nvshmem.UniqueID``, optional): A unique identifier used for UID-based initialization.
            Must be provided if `initializer_method` is "uid".
        - rank (int, optional): Rank of the calling process in the NVSHMEM job. Required for UID-based init.
        - nranks (int, optional): Total number of NVSHMEM ranks in the job. Required for UID-based init.
        - mpi_comm (``mpi4py.MPI.Comm``, optional): MPI communicator to use for MPI-based initialization.
            Defaults to ``MPI.COMM_WORLD`` if ``None`` and ``initializer_method`` is "mpi".
        - initializer_method (str): Specifies the initialization method. Must be either "mpi" or "uid".

    Raises:
        - NvshmemInvalid: If an invalid initialization method is provided, or required arguments
            for the selected method are missing or incorrect.

        - NvshmemError: If NVSHMEM fails to initialize using the specified method.

    Notes:
        - If using MPI-based init, ensure ``mpi4py`` is compiled against the same MPI
          distribution you're running with. A mismatch can result in undefined behavior.
        - For UID-based init, the user is responsible for distributing the `uid` to all
          processes and passing in the correct ``rank`` and ``nranks``.
        - UID-based init is useful for bootstrapping over non-MPI runtimes or custom transports.
        - Internally, this sets up a ``bindings.InitAttr()`` structure which is passed to
          the NVSHMEM host library.
        - This function uses the ``cuda.core`` Pathfinder module to locate and load the NVSHMEM host library.
          The documentation for this function, including the search order can be found at 
          https://nvidia.github.io/cuda-python/cuda-pathfinder/latest/generated/cuda.pathfinder.load_nvidia_dynamic_lib.html#cuda.pathfinder.load_nvidia_dynamic_lib

    Example:
        >>> from mpi4py import MPI
        >>> import nvshmem.core as nvshmem
        >>> nvshmem.init(mpi_comm=MPI.COMM_WORLD, initializer_method="mpi")
        # OR for UID mode
        >>> uid = nvshmem.get_unique_id()
        >>> nvshmem.init(uid=uid, rank=0, nranks=1, initializer_method="uid")
    """
    # If Device is None, that's ok.
    _cached_device["device"] = device

    # Load the nvshmemem host library
    # Must happen before we use bindings
    load_nvidia_dynamic_lib("nvshmem_host")

    attr = bindings.InitAttr()

    if initializer_method not in ["mpi", "uid", "emulated_mpi"]:
        raise NvshmemInvalid("Invalid init method requested")

    if initializer_method == "mpi":
        """
        The user requested MPI bootstrap. This means we internally use MPI to perform our bootstrap.
        It hasn't been done yet 
        """
        import sys
        print("Importing MPI", file=sys.stderr, flush=True)
        logger.debug("Importing MPI")
        _import_mpi()
        print(f"MPI4Py enabled: {_mpi4py_enabled}", file=sys.stderr, flush=True)
        logger.info(f"MPI4Py enabled after import: {_mpi4py_enabled}")
        # Step 1: Detect MPI_Comm size
        if not _mpi4py_enabled:
            raise NvshmemInvalid("MPI/MPI4Py not enabled")
        if mpi_comm is not None and not isinstance(mpi_comm, mpi.Comm):
            raise NvshmemInvalid("Invalid MPI communicator passed")

        # This has the effect of detecting the MPI distro
        # OMPI uses void *, MPICH family uses int
        # Note that as of today, MPI4Py defines this statically by compiling
        # User must make sure the MPI4Py they have is compiled for the MPI distro they are using.
        # Otherwise, weird errors will occur
        if mpi._sizeof(mpi.Comm) == ctypes.sizeof(ctypes.c_int):
            MPI_Comm_t = ctypes.c_int
        else:
            MPI_Comm_t = ctypes.c_void_p

        if mpi_comm is None:
            # If None, assume user wants COMM_WORLD
            mpi_comm = mpi.COMM_WORLD

        # Cast to the correct type for the MPI distro we found
        mpi_comm_val = MPI_Comm_t(mpi_comm.handle)
        # Get pointer
        mpi_comm_ptr = ctypes.pointer(mpi_comm_val)
        # Cast to void* for portability
        mpi_comm_addr = ctypes.cast(mpi_comm_ptr, ctypes.c_void_p).value

        attr.mpi_comm = mpi_comm_addr
        init_status = bindings.hostlib_init_attr(bindings.Flags.INIT_WITH_MPI_COMM, attr.ptr)

        # Cybind converts the success status code (0) to NoneType
        if init_status != None:
            raise NvshmemError("Failed to perform MPI based init")

    if initializer_method == "uid":
        """
        The user requested a UID bootstrap. That means they took it on themself to distribute the UID
        We assume they know what their doing and already broadcasted the UID to anyone who needs it
        
        The user should have gotten the UID from `nvshmem.core.get_unique_id()`. Only one rank calls this function.

        In future, if UID is None, MPI_comm is not None, and initializer method is "uid",
        we may want to provide a convenience functionality for using MPI internally to do the UID bootstrap.
        """
        if uid is None or rank is None or nranks is None or mpi_comm:
            raise NvshmemInvalid("Invalid parameters for UID based init {uid} {rank} {nranks}")

        # It's the caller's responsibility to pass in correct rank and nranks if they use UID bootstrap
        uid_status = bindings.set_attr_uniqueid_args(rank, nranks, uid.ptr, attr.ptr)
        status = bindings.hostlib_init_attr(bindings.Flags.INIT_WITH_UNIQUEID, attr.ptr)

        # Convert status to Exception
        # Cybind converts the success status code to NoneType
        if uid_status != None or status != None:
            raise NvshmemError("Failed to perform UID-based Init. status = {status}")

    if initializer_method == "emulated_mpi":
        """
        This method is an "emulated MPI init" where the user passes an MPI comm, and instead of passing it 
        into nvshmem_init directly, we perform a broadcast across it.

        This is useful for times when you don't want to recompile NVSHMEM for your MPI distro du-jour.
        """
        _import_mpi()
        if not _mpi4py_enabled:
            raise NvshmemInvalid("MPI4Py Required for emulated_mpi init")
        if mpi_comm is None:
            # If None, assume user wants COMM_WORLD
            mpi_comm = mpi.COMM_WORLD
        rank = mpi_comm.Get_rank()
        nranks = mpi_comm.Get_size()
        local_rank_per_node = rank % system.get_num_devices()

        # Create an empty uniqueid for all ranks
        uniqueid = get_unique_id(empty=True)
        if rank == 0:
            # Rank 0 gets a real uniqueid
            uniqueid = get_unique_id()

        # Broadcast UID to all ranks
        mpi_comm.Bcast(uniqueid._data.view(np.int8), root=0)

        # It's the caller's responsibility to pass in correct rank and nranks if they use UID bootstrap
        uid_status = bindings.set_attr_uniqueid_args(rank, nranks, uniqueid.ptr, attr.ptr)
        status = bindings.hostlib_init_attr(bindings.Flags.INIT_WITH_UNIQUEID, attr.ptr)

        # Convert status to Exception
        # Cybind converts the success status code to NoneType
        if uid_status != None or status != None:
            raise NvshmemError("Failed to perform emulated_mpi-based Init. status = {status}")

    log_level = os.environ.get("NVSHMEM_DEBUG")
    if log_level in ("INFO", "DEBUG"):
        _debug_mode = True

    if not log_level or log_level not in ["DEBUG", "INFO", "WARNING", "ERROR", None]:
        # Default to ERROR level
        log_level = "ERROR"

    # Now that we have PE set, re-initialize logging with mype None so that the logger picks up the correct one.
    utils._configure_logging(level=log_level)

    nvshmem.core._internal_tracking._is_initialized["status"] = InternalInitStatus.INITIALIZED


def finalize() -> None:
    """
    Finalize the NVSHMEM runtime.

    This function wraps the NVSHMEM finalization routine. It should be called after all 
    NVSHMEM operations are complete and before the application exits.

    Typically, this is called once per process to clean up NVSHMEM resources.

    Raises:
        ``NvshmemError``: If the NVSHMEM finalization fails.

    Example:
        >>> nvshmem.core.finalize()
    """
    logger.debug("nvshmem_finalize() called")
    memory._free_all_buffers()

    fini_status = bindings.hostlib_finalize()
    # Cybind converts the success status code to NoneType
    if fini_status != None:
        raise NvshmemError("Failed to finalize Hostlib")

    nvshmem.core._internal_tracking._is_initialized["status"] = InternalInitStatus.DE_INITIALIZED


def module_init(mod: NvshmemKernelObject) -> None:
    """
    Initialize the CUmodule instance backing the compiled object binary. The instance is of cuda.core.ObjectCode type.

    Typically, this is called once per uniquely compiled object, prior to launching the kernel. Internally, it can manage multiple init/finalize via refcounting.

    The same object passed in here must be passed into the corresponding module_finalize call. This function modifies the object in place.

    Args:
        - mod (NvshmemKernelObject): The same object passed in here must be passed into the corresponding module_finalize call. This function modifies the object in place.

    Raises:
        NvshmemError: If the NVSHMEM module initialization fails.

    Example:
        >>> nvshmem.core.module_init(mod)
    """
    if mod.handle is None:
        raise NvshmemInvalid("Invalid module type passed in")
    status = bindings.cumodule_init(int(mod.handle))
    print(f"CUmodule init status: {status}")
    if status is not None and status != 0:
        raise NvshmemError("Failed to initialize CUmodule for NVSHMEM")


def module_finalize(mod: NvshmemKernelObject) -> None:
    """
    Finalize the CUmodule instance backing the compiled object binary. The instance is of cuda.core.ObjectCode type.

    This is called once per uniquely compiled object, post teardown of the kernel.

    The user should pass the same NvshmemKernelObject that was passed to module_init into this function.

    Args:
        - mod (NvshmemKernelObject): The module that NVSHMEM will finalize

    Raises:
        NvshmemError: If the NVSHMEM module finalization fails.

    Example:
        >>> nvshmem.core.module_finalize(mod)
    """
    # At init time, we stored the handle of the loaded module
    if mod.finalize_handle is None:
        raise NvshmemInvalid("Module not initialized")
    status = bindings.cumodule_finalize(int(mod.handle))
    if status is not None and status != 0:
        raise NvshmemError("Failed to finalize CUmodule for NVSHMEM")


def library_init(lib: NvshmemKernelObject) -> None:
    """
    Initialize the CUmodule instance backing the compiled object binary. The instance is of cuda.core.ObjectCode type.

    Typically, this is called once per uniquely compiled object, prior to launching the kernel. Internally, it can manage multiple init/finalize via refcounting.

    The same object passed in here must be passed into the corresponding library_finalize call. This function modifies the object in place.

    Args:
        - lib (NvshmemKernelObject): The library that NVSHMEM will initialize
    
    Raises:
        NvshmemError: If the NVSHMEM module initialization fails.

            The same object passed in here must be passed into the corresponding module_finalize call. This function modifies the object in place.

    Example:
        >>> nvshmem.core.module_init(mod)
    """
    if lib.handle is None:
        raise NvshmemInvalid("Invalid library type passed in")
    status = bindings.culibrary_init(lib.handle)
    if status is not None and status != 0:
        raise NvshmemError("Failed to initialize CULibrary for NVSHMEM")


def library_finalize(lib: NvshmemKernelObject) -> None:
    """
    Finalize the CUmodule instance backing the compiled object binary. The instance is of cuda.core.ObjectCode type.

    Typically, this is called once per uniquely compiled object, post teardown of the kernel. Internally, it can manage multiple init/finalize via refcounting.
    
    Args:
        - lib (NvshmemKernelObject): The library that NVSHMEM will finalize

    The user should pass the same NvshmemKernelObject that was passed to library_init into this function.

    Raises:
        NvshmemError: If the NVSHMEM module finalization fails.

    Example:
        >>> nvshmem.core.module_finalize(mod)
    """
    # At init time, we stored the handle of the loaded module
    if lib.handle is None:
        raise NvshmemInvalid("Library not initialized")
    status = bindings.culibrary_finalize(int(lib.handle))
    if status is not None and status != 0:
        raise NvshmemError("Failed to finalize CULibrary for NVSHMEM")


def _normalize_device_library_arch(arch: object) -> str:
    if isinstance(arch, (tuple, list)) and len(arch) >= 2:
        return f"{int(arch[0])}{int(arch[1])}"

    major = getattr(arch, "major", None)
    minor = getattr(arch, "minor", None)
    if major is not None and minor is not None:
        return f"{int(major)}{int(minor)}"

    arch_str = str(arch).lower().strip()
    if arch_str.startswith("compute_"):
        arch_str = arch_str[8:]
    if arch_str.startswith("sm_"):
        arch_str = arch_str[3:]
    elif arch_str.startswith("sm"):
        arch_str = arch_str[2:]
    arch_str = arch_str.replace(".", "").rstrip("a")
    if arch_str.isdigit():
        return arch_str

    parts = re.findall(r"\d+", arch_str)
    if len(parts) == 2:
        return f"{int(parts[0])}{int(parts[1])}"

    raise ValueError(f"Could not normalize device library arch: {arch}")


def _current_device_library_arch() -> Optional[str]:
    try:
        device = Device()
    except Exception:
        return None

    for attr in ("compute_capability", "arch"):
        try:
            value = getattr(device, attr)
        except Exception:
            continue
        if value is None:
            continue
        try:
            return _normalize_device_library_arch(value)
        except Exception:
            continue
    return None


def find_device_bitcode_library(
    *,
    language: DeviceLibLanguage = DeviceLibLanguage.LLVM_BITCODE,
    arch: Optional[Union[int, str]] = None,
) -> str:
    """
    Find the path to an NVSHMEM device library.

    Searches for the library in ../lib/ relative to the NVSHMEM header directory.

    Args:
        language: Library format to search for. ``DeviceLibLanguage.LLVM_BITCODE``
            returns a ``.bc`` file (for use with ``llvm-link``).
            ``DeviceLibLanguage.LTOIR`` returns the installed
            ``libnvshmem_device.ltoir.fatbin`` entry point. On CUDA 13+ this is
            a multi-arch fatbin; on CUDA 12.8-12.x the file keeps that suffix for
            layout consistency but is a renamed single-arch ``sm_90`` LTOIR image.
            CUDA 12 consumers that need per-arch LTOIR should use the installed
            ``libnvshmem_device_sm_<arch>.ltoir`` file directly.
            ``DeviceLibLanguage.STATIC`` returns the ``.a`` static archive
            (for use with standard device linking).
        arch: SM architecture number (e.g. ``"90"``, ``"sm_90"``,
            ``"80"``, ``"100"``). By default, LLVM_BITCODE lookup prefers
            the current device's per-arch file when available, then falls back
            to the backward-compatible ``libnvshmem_device.bc`` entry point.
            For LLVM_BITCODE, selects the requested per-arch file
            (e.g. ``libnvshmem_device_sm_90.bc``) when this value is set.
            For LTOIR and STATIC, this helper returns the install-layout entry
            point regardless of this value.

    Returns:
        The path to the device library file. (a string)
    """
    header_path = find_nvidia_header_directory("nvshmem")
    if not header_path:
        raise NvshmemInvalid("NVSHMEM headers not found. Cannot find the device bitcode library.")

    # Search in ../lib/ relative to the header path
    # TODO: Switch to cuda.pathfinder when it supports bitcode libraries (https://github.com/NVIDIA/cuda-python/issues/1421)
    header_path_obj = Path(header_path)
    if language == DeviceLibLanguage.LTOIR:
        filename = "libnvshmem_device.ltoir.fatbin"
    elif language == DeviceLibLanguage.STATIC:
        filename = "libnvshmem_device.a"
    else:
        if arch is None:
            candidates = []
            current_arch = _current_device_library_arch()
            if current_arch is not None:
                candidates.append(f"libnvshmem_device_sm_{current_arch}.bc")
            candidates.append("libnvshmem_device.bc")
        else:
            arch = _normalize_device_library_arch(arch)
            candidates = [f"libnvshmem_device_sm_{arch}.bc"]

        for filename in candidates:
            lib_path = header_path_obj.parent / "lib" / filename
            if lib_path.exists():
                return str(lib_path)
        raise NvshmemInvalid(f"NVSHMEM device library not found. Tried: {', '.join(candidates)}")

    lib_path = header_path_obj.parent / "lib" / filename
    if not lib_path.exists():
        raise NvshmemInvalid(f"NVSHMEM device library not found at {lib_path}")
    return str(lib_path)
