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

cimport cpython

from libc.stddef cimport size_t
from cuda.bindings cimport cydriver
from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle
from cuda.core._memory._memory_pool cimport _MemPool
from cuda.core._stream cimport Stream, Stream_accept
from cuda.core._resource_handles cimport (
    DevicePtrHandle,
    create_fd_handle,
    create_mempool_handle_ipc,
    deviceptr_import_ipc,
    get_last_error,
    as_cu,
    as_intptr,
    as_py,
)

from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
from cuda.core._utils.cuda_utils import check_multiprocessing_start_method

import multiprocessing
import os
import platform
import uuid
import weakref

__all__ = []


cdef object registry = weakref.WeakValueDictionary()


cdef cydriver.CUmemAllocationHandleType IPC_HANDLE_TYPE =                       \
    cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR \
    if platform.system() == "Linux" else                                        \
    cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE

cdef is_supported():
    return IPC_HANDLE_TYPE != cydriver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE


cdef class IPCDataForBuffer:
    """Data members related to sharing memory buffers via IPC."""
    def __cinit__(self, IPCBufferDescriptor ipc_descriptor, bint is_mapped) -> None:
        self._ipc_descriptor = ipc_descriptor
        self._is_mapped = is_mapped

    @property
    def ipc_descriptor(self) -> IPCBufferDescriptor:
        return self._ipc_descriptor

    @property
    def is_mapped(self) -> bool:
        return self._is_mapped


cdef class IPCDataForMR:
    """Data members related to sharing memory resources via IPC."""
    def __cinit__(self, IPCAllocationHandle alloc_handle, bint is_mapped) -> None:
        self._alloc_handle = alloc_handle
        self._is_mapped = is_mapped

    @property
    def alloc_handle(self) -> IPCAllocationHandle:
        return self._alloc_handle

    @property
    def is_mapped(self) -> bool:
        return self._is_mapped

    @property
    def uuid(self) -> uuid.UUID | None:
        return getattr(self._alloc_handle, 'uuid', None)


cdef class IPCBufferDescriptor:
    """Serializable object describing a buffer that can be shared between processes.

    Note
    ----
    The payload and ``size`` fields are controlled by the exporting peer.
    Receivers must treat them as untrusted and import only through
    :meth:`Buffer.from_ipc_descriptor`.
    """

    def __init__(self, *arg, **kwargs) -> None:
        raise RuntimeError("IPCBufferDescriptor objects cannot be instantiated directly. Please use MemoryResource APIs.")

    @staticmethod
    def _init(reserved: bytes, size: int) -> IPCBufferDescriptor:
        cdef IPCBufferDescriptor self = IPCBufferDescriptor.__new__(IPCBufferDescriptor)
        self._payload = reserved
        self._size = size
        return self

    def __reduce__(self) -> tuple[object, ...]:
        return IPCBufferDescriptor._init, (self._payload, self._size)

    @property
    def size(self) -> int:
        return self._size

    cdef const void* payload_ptr(self) noexcept:
        """Return the payload as a const void* for C API calls."""
        return <const void*><const char*>(self._payload)


cdef class IPCAllocationHandle:
    """Shareable handle to an IPC-enabled device memory pool."""

    def __init__(self, *arg, **kwargs) -> None:
        raise RuntimeError("IPCAllocationHandle objects cannot be instantiated directly. Please use MemoryResource APIs.")

    @classmethod
    def _init(cls, handle: int, uuid: uuid.UUID | None) -> IPCAllocationHandle:  # no-cython-lint
        cdef IPCAllocationHandle self = IPCAllocationHandle.__new__(cls)
        if handle < 0:
            raise ValueError(f"Invalid allocation handle (fd) {handle}: must be non-negative")
        self._h_fd = create_fd_handle(handle)
        self._uuid = uuid
        return self

    cpdef close(self):
        """Close the handle."""
        self._h_fd.reset()

    def __int__(self) -> int:
        if not self._h_fd or as_intptr(self._h_fd) < 0:
            raise ValueError(
                f"Cannot convert IPCAllocationHandle to int: the handle (id={id(self)}) is closed."
            )
        return as_py(self._h_fd)

    @property
    def handle(self) -> int:
        return as_py(self._h_fd)

    @property
    def uuid(self) -> uuid.UUID:
        return self._uuid


def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]:
    check_multiprocessing_start_method()
    df = multiprocessing.reduction.DupFd(alloc_handle.handle)
    return _reconstruct_allocation_handle, (type(alloc_handle), df, alloc_handle.uuid)


def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None) -> IPCAllocationHandle:  # no-cython-lint
    return cls._init(df.detach(), uuid)


multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle)


# Buffer IPC Implementation
# -------------------------
cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self):
    if not self.memory_resource.is_ipc_enabled:
        raise RuntimeError("Memory resource is not IPC-enabled")
    cdef cydriver.CUmemPoolPtrExportData data
    with nogil:
        HANDLE_RETURN(
            cydriver.cuMemPoolExportPointer(&data, as_cu(self._h_ptr))
        )
    cdef bytes data_b = cpython.PyBytes_FromStringAndSize(
        <char*>(data.reserved), sizeof(data.reserved)
    )
    return IPCBufferDescriptor._init(data_b, self.size)

cdef Buffer Buffer_from_ipc_descriptor(
    cls, _MemPool mr, IPCBufferDescriptor ipc_descriptor, stream
):
    """Import a buffer that was exported from another process."""
    if not mr.is_ipc_enabled:
        raise RuntimeError("Memory resource is not IPC-enabled")
    cdef size_t payload_size = len(ipc_descriptor._payload)
    cdef size_t expected_size = sizeof(cydriver.CUmemPoolPtrExportData)
    if payload_size < expected_size:
        raise ValueError(
            f"IPC buffer descriptor payload is {payload_size} bytes; "
            f"expected at least {expected_size}"
        )
    cdef Stream s = Stream_accept(stream)
    cdef DevicePtrHandle h_ptr = deviceptr_import_ipc(
        mr._h_pool,
        ipc_descriptor.payload_ptr(),
        s._h_stream
    )
    if not h_ptr:
        HANDLE_RETURN(get_last_error())
    cdef size_t mapped_size = 0
    cdef size_t claimed_size = ipc_descriptor.size
    with nogil:
        HANDLE_RETURN(cydriver.cuPointerGetAttribute(
            &mapped_size,
            cydriver.CU_POINTER_ATTRIBUTE_RANGE_SIZE,
            as_cu(h_ptr)))
    if claimed_size > mapped_size:
        h_ptr.reset()
        raise ValueError(
            f"IPC buffer descriptor size ({claimed_size}) exceeds "
            f"mapped allocation extent ({mapped_size} bytes)"
        )
    return Buffer_from_deviceptr_handle(h_ptr, claimed_size, mr, ipc_descriptor)


# _MemPool IPC Implementation
# ---------------------------

cdef _MemPool MP_from_allocation_handle(cls, alloc_handle):
    # Quick exit for registry hits.
    uuid = getattr(alloc_handle, 'uuid', None)  # no-cython-lint
    mr = registry.get(uuid)
    if mr is not None:
        if not isinstance(mr, cls):
            raise TypeError(
                f"Registry contains a {type(mr).__name__} for uuid "
                f"{uuid}, but {cls.__name__} was requested")
        return mr

    # Ensure we have an allocation handle. Duplicate the file descriptor, if
    # necessary.
    if isinstance(alloc_handle, int):
        fd = os.dup(alloc_handle)
        try:
            alloc_handle = IPCAllocationHandle._init(fd, None)
        except:
            os.close(fd)
            raise

    # Construct a new mempool.
    cdef _MemPool self = <_MemPool>(cls.__new__(cls))
    self._mempool_owned = True
    cdef int ipc_fd = int(alloc_handle)
    self._h_pool = create_mempool_handle_ipc(ipc_fd, IPC_HANDLE_TYPE)
    if not self._h_pool:
        HANDLE_RETURN(get_last_error())
        raise RuntimeError(
            f"Failed to import {cls.__name__} from an allocation handle: "
            "cuda-core returned an empty memory pool handle without recording a CUDA error. "
            "This is an internal cuda-core error; please report it with your CUDA driver, "
            "CUDA Toolkit, and cuda-python versions."
        )
    self._ipc_data = IPCDataForMR(alloc_handle, True)

    # Register it.
    if uuid is not None:
        registered = self.register(uuid)
        assert registered is self

    return self


cdef _MemPool MP_from_registry(uuid):
    try:
        return registry[uuid]
    except KeyError:
        raise RuntimeError(f"Memory resource {uuid} was not found") from None


cdef _MemPool MP_register(_MemPool self, uuid):
    existing = registry.get(uuid)
    if existing is not None:
        return existing
    assert self.uuid is None or self.uuid == uuid
    registry[uuid] = self
    self._ipc_data._alloc_handle._uuid = uuid
    return self


cdef IPCAllocationHandle MP_export_mempool(_MemPool self):
    # Note: This is Linux only (int for file descriptor)
    cdef int fd
    with nogil:
        HANDLE_RETURN(cydriver.cuMemPoolExportToShareableHandle(
            &fd, as_cu(self._h_pool), IPC_HANDLE_TYPE, 0)
        )
    try:
        return IPCAllocationHandle._init(fd, uuid.uuid4())
    except:
        os.close(fd)
        raise
