# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""
These are the Python datatypes for NVSHMEM
"""
import uuid
import logging
from typing import Union
from enum import Enum, IntEnum

from cuda.core import MemoryResource, Buffer
from cuda.core import Device, system, ObjectCode
from cuda.core import Stream
from cuda.bindings.utils import get_cuda_native_handle
import cuda.bindings.driver

from nvshmem.bindings import malloc, free, ptr, mc_ptr, Team_id, buffer_register_symmetric, buffer_unregister_symmetric
from nvshmem.core._internal_tracking import _is_initialized, InternalInitStatus

logger = logging.getLogger("nvshmem")

__all__ = [
    "Version", "NvshmemInvalid", "NvshmemError", "NvshmemResource", "NvshmemStream", "NvshmemStreamsType",
    "MutableEnum", "Teams", "BufferTypes", "NvshmemKernelObject"
]

try:
    import torch
    from torch.cuda import Stream as torch_stream
    _torch_enabled = True
except:
    torch_stream = None
"""
Mutable Enum to support team creation
"""


class MutableEnum:

    def __init__(self, **kwargs):
        self._members = {}
        for key, value in kwargs.items():
            self._members[key] = value

    @classmethod
    def from_enum(cls, enum_cls: Enum):
        # Extract members from the standard enum class and create a MutableEnum instance
        members = {member.name: member.value for member in enum_cls}
        return cls(**members)

    def add(self, key, value):
        if key in self._members:
            raise ValueError(f"Member '{key}' already exists.")
        self._members[key] = value

    def remove(self, key):
        if key not in self._members:
            raise ValueError(f"Member '{key}' does not exist.")
        del self._members[key]

    def remove_by_value(self, value):
        """
        Remove a member by value.
        This is useful when we want to remove a team by its handle.
        Calling this function on a value notin the Enum is idempotent.

        Args:
            - value (int): The value of the member to remove.

        Returns:
            None

        Raises:
        """
        for key, val in self._members.items():
            if val == value:
                del self._members[key]
                return

    def __getattr__(self, item):
        if item in self._members:
            return self._members[item]
        raise AttributeError(f"No such member: {item}")

    def __getitem__(self, item):
        if item in self._members:
            return self._members[item]
        raise KeyError(f"No such member: {item}")

    def __contains__(self, item):
        return item in self._members

    def __repr__(self):
        members = ", ".join(f"{k}={v}" for k, v in self._members.items())
        return f"<MutableEnum {members}>"

    def keys(self):
        return self._members.keys()

    def values(self):
        return self._members.values()

    def items(self):
        return self._members.items()

    def __iter__(self):
        return iter(self._members)


"""
IntEnum which matches 1:1 with ``nvshmem_team_id_t``

This is a special Mutable enum so we can add elements to it later after a team was created.
"""
Teams = MutableEnum.from_enum(Team_id)


class NvshmemKernelObject:

    def __new__(self, *args, **kwargs):
        raise RuntimeError("NvshmemKernelObject objects cannot be instantiated directly. "
                           "Please use NvshmemKernelObject APIs (from_cubin, from_handle)")

    @classmethod
    def _init(cls, handle: int):
        self = super().__new__(cls)
        self.handle = handle
        return self

    @staticmethod
    def from_handle(handle: int = None):
        """
        Create a NvshmemKernelObject from a handle
        """
        return NvshmemKernelObject._init(handle=int(handle))

    @staticmethod
    def from_obj(obj: ObjectCode):
        """
        Create a NvshmemKernelObject from a ObjectCode
        """
        return NvshmemKernelObject._init(handle=get_cuda_native_handle(obj.handle))

    @staticmethod
    def from_obj_bytes(obj_bytes: bytes):
        """
        Create a NvshmemKernelObject from a bytes-string of the assembled CUDA object.
        """
        obj = ObjectCode.from_cubin(obj_bytes)
        return NvshmemKernelObject._init(handle=get_cuda_native_handle(obj.handle))

    @staticmethod
    def from_triton(triton_kernel: "JITFunction"):
        """
        Create a NvshmemKernelObject from a Triton kernel compiled with JITFunction
        """
        return NvshmemKernelObject._init(handle=int(triton_kernel.handle))

    # TODO: convert to class when cuda.bindings.Kernel is imported
    @staticmethod
    def from_cuda_bindings(cuda_kernel: "cuda.bindings.driver.CUmodule"):
        """
        Create a NvshmemKernelObject from a CUDA kernel
        """
        return NvshmemKernelObject._init(handle=get_cuda_native_handle(cuda_kernel))


"""
Version class
"""


class Version:

    def __init__(self, openshmem_spec_version="", nvshmem4py_version="", libnvshmem_version=""):
        self.openshmem_spec_version = openshmem_spec_version
        self.nvshmem4py_version = nvshmem4py_version
        self.libnvshmem_version = libnvshmem_version

    def __repr__(self):
        return f"""
NVSHMEM4Py Library:
    NVSHMEM Library version: {self.libnvshmem_version}
    OpenShmem Spec version: {self.openshmem_spec_version}
    NVSHMEM4Py version: {self.nvshmem4py_version}
"""


"""
Wrapper class for cuda.core StreamInteroperability protocol
"""


class NvshmemStream:

    def __init__(self, pt_stream):
        self.pt_stream = pt_stream
        self.handle = pt_stream.cuda_stream

    def __cuda_stream__(self):
        stream_id = self.pt_stream.cuda_stream
        return (0, stream_id)  # Return format required by CUDA Python


NvshmemStreamsType = Union[torch_stream, NvshmemStream, Stream]
"""
Exceptions
"""


class NvshmemInvalid(Exception):

    def __init__(self, msg):
        self.msg = msg

    def __repr__(self):
        return f"<NvshmemInvalid: {self.msg}>"


class NvshmemError(Exception):

    def __init__(self, msg):
        self.msg = msg

    def __repr__(self):
        return f"<NvshmemError: {self.msg}>"


class NvshmemWarning(UserWarning):

    def __init__(self, msg):
        self.msg = msg

    def __repr__(self):
        return f"<NvshmemWarning: {self.msg}>"


"""
Buffer types for NvshmemResource
"""


class BufferTypes(IntEnum):
    NORMAL = 0  # Normal buffer
    PEER = 1  # Peer buffer (nvshmem_ptr)
    MULTIMEM = 2  # Multicast buffer (nvshmemx_mc_ptr)
    EXTERNAL = 3  # External buffer (nvshmemx_buffer_register)


"""
Memory Resource
"""


class NvshmemResource(MemoryResource):
    """
    A memory resource that uses NVSHMEM to allocate device memory.

    This class implements the MemoryResource interface and allocates memory using
    ``nvshmem_malloc``. It supports device-accessible memory but does not allow host access.

    Attributes:
        device (Device): The CUDA device associated with this resource.
    """

    def __init__(self, device):
        """
        Initialize the NVSHMEM memory resource for a specific device.

        Args:
            device (Device): The CUDA device object on which memory will be allocated.
        """
        self.device = device
        """
        Map of symmetric heap pointers to nvshmem.core objects
        Keys: ptr, values: 
            {
            # A reference count
            "ref_count": <int>,
            # An object of type NvshmemResource
            "device": <int>,
            # An object of type cuda.core.Buffer
            "buffer": <Buffer>,
            # Different buffers (peer, mc) have different needs
            "type": Enum,
            # Used to raise an exception when the GC reaches here without free() getting called
            "freed": Bool,
            # Used to track whether the user expects us to prevent them from experiencing deadlocks
            "released": Bool,
            # Used to track whether the user expects us to raise an exception when the buffer is not tracked or has already been freed
            # True -> raise an exception
            # False -> print an error message
            "except_on_del": Bool
            }
        """
        self._mem_references = {}

    def allocate(self, size: int, stream: Stream = None, release=False, except_on_del=True) -> Buffer:
        """
        Allocate memory on the device using NVSHMEM.

        Args:
            - size (int): Number of bytes to allocate.
            - stream (optional): CUDA stream for allocation context (not used here).
            - release (bool, optional): do not track the buffer if True
            - except_on_del (bool, optional): If True, raise an exception if the buffer is not tracked or has already been freed.
                If False, print an error message.

        Returns:
            ``Buffer``: A buffer object wrapping the allocated device memory.

        Raises:
            ``NvshmemError``: If the allocation fails.
        """

        ptr = malloc(size)
        if not ptr or ptr == 0:
            raise NvshmemError(f"Failed to allocate memory of bytes {size}")
        r_buf = Buffer.from_handle(ptr=ptr, size=size, mr=self)
        logger.debug(f"Created Buffer on resource {self} at address {ptr} with size {size} on stream {stream}")

        if not release:
            buf_ref = r_buf
        else:
            # If we're not holding references, create our own shadow-buffer
            buf_ref = Buffer.from_handle(ptr=r_buf.handle, size=r_buf.size, mr=self)
        self._mem_references[ptr] = {
            "ref_count": 1,
            "resource": self,
            "buffer": buf_ref,
            "type": BufferTypes.NORMAL,
            "freed": False,
            "released": release,
            "except_on_del": except_on_del,
            "child": set()
        }

        return r_buf

    def deallocate(self, ptr: int, size: int, stream: NvshmemStreamsType = None) -> None:
        """
        Placeholder method for deallocation.

        This function will be called when a Buffer has ``.close()`` called.
        ``.close()`` must be called explicitly

        Args:
            - ptr (int): Pointer to the memory block (ignored).
            - size (int): Size of the memory block in bytes (ignored).
            - stream (optional): CUDA stream (ignored).

        Returns:
            None
        """
        # Extract info
        logger.debug(f"Free called on buffer with address {ptr}")
        if not _is_initialized["status"] == InternalInitStatus.INITIALIZED:
            logger.warning("NVSHMEM Library is not initialized. Cannot free buffer")
            return

        if not hasattr(self, "_mem_references"):
            logger.info("Cannot free buffer. NVSHMEM resource is being destroyed")
            return
        if self._mem_references.get(ptr) is None:
            logger.debug("Freed a buffer that is not tracked")
            return
        released = self._mem_references[ptr].get("released", False)

        # If someone got here without calling free(), we have to except
        if self._mem_references[ptr]["type"] == BufferTypes.NORMAL and not self._mem_references[ptr]["freed"]:
            if self._mem_references[ptr].get("except_on_del", True):
                raise NvshmemError(f'Buffer {self._mem_references[ptr]["buffer"]} freed implicitly.')
            else:
                if not self._mem_references[ptr].get("warned", False):
                    logger.error(f'Buffer {self._mem_references[ptr]["buffer"]} freed implicitly.')
                    self._mem_references[ptr]["warned"] = True
                return

        # remove the reference
        # We keep the references around to legalize freed-ness.
        # If nvshmem_malloc returns a new ptr, we will end up creating a new entry with refcount 1
        # and freed False
        if self._mem_references[ptr]["ref_count"] > 0:
            self._mem_references[ptr]["ref_count"] -= 1
        elif self._mem_references[ptr]["ref_count"] == 0:
            # The counter is already 0, so we must have already freed the pointer. Just return.
            logger.debug(f"Ref count on {ptr} is already 0. Already freed.")
            return
        logger.debug(
            f"New ref count on {self._mem_references[ptr]['type'].name} buf {ptr} {self._mem_references[ptr]['ref_count'] }"
        )
        # If this was the last reference to that pointer, free the pointer
        # The MR itself has a ref_count, but we want to free only when the last call is made.
        # Leave this as if ( == 1) and if it's 0, we will delete the reference
        if self._mem_references[ptr]["ref_count"] == 0:
            # If the buffer is a peer buffer, don't do anything
            # except delete it from the tracker
            # NVShmem handles these internally.
            if self._mem_references[ptr]["type"] == BufferTypes.NORMAL:
                free(ptr)
                # Free any peer/multicast children tracked for this parent.
                child_ptrs = self._mem_references[ptr]["child"]
                for child_ptr in child_ptrs:
                    child_entry = self._mem_references.get(child_ptr)
                    if child_entry is None:
                        continue
                    if child_entry.get("type") not in (BufferTypes.PEER, BufferTypes.MULTIMEM):
                        continue
                    child_buffer = child_entry.get("buffer", None)
                    if child_buffer is not None:
                        self._mem_references[child_ptr]["freed"] = True
                        del self._mem_references[child_ptr]["buffer"]
                self._mem_references[ptr]["child"].clear()
                logger.debug(f"Freed buffer at address {ptr}")
            else:
                logger.debug("free() requested on a peer buffer. Not calling free()")
                self._mem_references[ptr]["freed"] = True
                buf_entry = self._mem_references[ptr].get("buffer", None)
                if buf_entry is not None:
                    del self._mem_references[ptr]["buffer"]

    def get_peer_buffer(self, buffer: Buffer, pe: int) -> Buffer:

        # This should be the pointer on the calling PE
        # None or raising an exception is the failing case
        parent_ptr = buffer.handle
        result = ptr(parent_ptr, pe)

        if not result or result == 0:
            raise NvshmemError("Failed to retrieve peer buffer")

        entry = self._mem_references.get(result, None)
        if entry is not None and not entry["freed"]:
            # Someone already called get_peer_buffer on the Buffer
            # Increase ref count and return existing buffer
            self._mem_references[result]["ref_count"] += 1
            logger.debug(
                f"Found already tracked peer buffer with address {result}. Returning it. Ref count {self._mem_references[result]['ref_count']}"
            )
            return self._mem_references[result]["buffer"]

        logger.debug(f"Did not find peer buffer with address {result}. Creating a new one.")

        # This Buffer doesn't need to go through any .allocate() calls, since we know the pointer is valid
        r_buf = Buffer.from_handle(ptr=result, size=buffer.size, mr=self)

        self._mem_references[result] = {
            "ref_count": 1,
            "resource": self,
            "buffer": r_buf,
            "type": BufferTypes.PEER,
            "freed": False,
            "parent": parent_ptr,
            "released": False,
            "except_on_del": False
        }
        parent_entry = self._mem_references[parent_ptr]
        parent_entry["child"].add(result)
        return r_buf

    def get_mc_buffer(self, team: Teams, buffer: Buffer) -> Buffer:

        # This should be the pointer on the calling PE
        # None or raising an exception is the failing case
        parent_ptr = buffer.handle
        result = mc_ptr(team, parent_ptr)
        if not result or result == 0:
            raise NvshmemError("Failed to retrieve multicast buffer")

        entry = self._mem_references.get(result, None)
        if entry is not None and not entry["freed"]:
            # Someone already called get_peer_buffer on the Buffer
            # Increase ref count and return existing buffer
            self._mem_references[result]["ref_count"] += 1
            logger.debug(
                f"Found already tracked MC buffer with address {result}. Returning it. Ref count {self._mem_references[result]['ref_count']}"
            )
            return self._mem_references[result]["buffer"]

        logger.debug(f"Did not find MC buffer with address {result}. Creating a new one.")

        # This Buffer doesn't need to go through any .allocate() calls, since we know the pointer is valid
        r_buf = Buffer.from_handle(ptr=result, size=buffer.size, mr=self)
        self._mem_references[result] = {
            "ref_count": 1,
            "resource": self,
            "buffer": r_buf,
            "type": BufferTypes.MULTIMEM,
            "freed": False,
            "parent": parent_ptr,
            "released": False,
            "except_on_del": False
        }
        parent_entry = self._mem_references[parent_ptr]
        parent_entry["child"].add(result)
        return r_buf

    def set_freed(self, buffer: Buffer) -> None:
        ptr = buffer.handle
        if self._mem_references.get(ptr) is None:
            raise NvshmemError("Freed a buffer that is not tracked")
        if self._mem_references[ptr]['type'] != BufferTypes.NORMAL:
            return
        self._mem_references[ptr]["freed"] = True

    def register_external_buffer(self, buffer: Buffer, call_register: bool = True) -> Buffer:
        """
        Register an external buffer with NVSHMEM.

        Args:
            - buffer (Buffer): The buffer to register.
            - call_register (bool, optional): Whether to call the register function on the buffer.
                                              This should be False if the buffer is already registered
                                              or was allocated via nvshmem_malloc outside of NVSHMEM4Py
        Returns:
            Buffer: A buffer object wrapping the registered external buffer.
            Users must pass this buffer to NVSHMEM operations, rather than the original buffer.
        """
        ptr = int(buffer.handle)
        if self._mem_references.get(ptr) is not None:
            if self._mem_references[ptr]['type'] != BufferTypes.EXTERNAL:
                raise NvshmemError(
                    "Tried to register an external buffer that is already tracked as a different type of buffer")
            self._mem_references[ptr]["ref_count"] += 1
            logger.warning(
                f"Found already tracked external buffer with address {ptr}. Returning it. Ref count {self._mem_references[ptr]['ref_count']}"
            )
            return self._mem_references[ptr].get("buffer", None)
        if call_register:
            registered_ptr = buffer_register_symmetric(int(ptr), int(buffer.size), 0)
            if registered_ptr is None:
                raise NvshmemError("Failed to register external buffer")
            registered_ptr = int(registered_ptr)
            logger.debug(f"Registered external buffer with address {ptr}.")
            new_buf = Buffer.from_handle(ptr=registered_ptr, size=buffer.size, mr=self)
            logger.debug(f"Created new buffer with address {new_buf.handle}.")
        else:
            registered_ptr = int(ptr)
            new_buf = buffer
        self._mem_references[int(registered_ptr)] = {
            "ref_count": 1,
            "resource": self,
            "buffer": new_buf,
            "type": BufferTypes.EXTERNAL,
            "freed": False,
            "released": not call_register
        }
        logger.debug(
            f"Registered external buffer with address {registered_ptr}. Ref count {self._mem_references[registered_ptr]['ref_count']}"
        )
        return new_buf

    def unregister_external_buffer(self, buffer: Buffer) -> None:
        ptr = int(buffer.handle)
        if self._mem_references.get(ptr) is None:
            logger.warning(
                f"Tried to unregister an external buffer that is not tracked with address {ptr}. Make sure you pass the buffer that NVSHMEM4Py returned to you"
            )
            raise NvshmemError("Tried to unregister an external buffer that is not tracked")
        if self._mem_references[ptr]['type'] != BufferTypes.EXTERNAL:
            raise NvshmemError("Tried to unregister an external buffer that is not external")
        del self._mem_references[ptr]["buffer"]
        self._mem_references[ptr]["ref_count"] = 0
        self._mem_references[ptr]["freed"] = True
        if not self._mem_references[ptr]["released"]:
            buffer_unregister_symmetric(ptr, buffer.size)

        self._mem_references[int(ptr)]["freed"] = True
        logger.debug(
            f"Unregistered external buffer with address {ptr}. Ref count {self._mem_references[ptr]['ref_count']}")

    @property
    def is_device_accessible(self) -> bool:
        """
        Indicates whether the allocated memory is accessible from the device.

        Returns:
            bool: Always True for NVSHMEM memory.
        """
        return True

    @property
    def is_host_accessible(self) -> bool:
        """
        Indicates whether the allocated memory is accessible from the host.

        Returns:
            bool: Always False for NVSHMEM memory.
        """
        return False

    @property
    def device_id(self) -> int:
        """
        Get the device ID associated with this memory resource.

        Returns:
            int: CUDA device ID.
        """
        return self.device.device_id

    def __repr__(self) -> str:
        """
        Return a string representation of the NvshmemResource.

        Returns:
            str: A string describing the object
        """
        return f"<NvshmemResource device={self.device}>"
