# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
# See License.txt for license information
#
# This code was automatically generated from NVSHMEM $version_span. 
# Modify it directly at your own risk.

cimport cython  # NOQA
from cpython cimport buffer as _buffer
from cpython.memoryview cimport PyMemoryView_FromMemory
from libc.stdint cimport (
    int8_t,  uint8_t,
    int16_t, uint16_t,
    int32_t, uint32_t,
    int64_t, uint64_t,
    intptr_t, uintptr_t
)


from enum import IntEnum as _IntEnum

import numpy as _numpy


###############################################################################
# POD
###############################################################################

$pod_defs

cdef class UniqueId(uniqueid): pass

# POD wrapper for nvshmemx_init_attr_t. cybind can't generate this automatically
# because it doesn't fully support nested structs (https://gitlab-master.nvidia.com/leof/cybind/-/issues/67).
# The nested structure is made opaque.
# TODO: remove this once cybind supports nested structs.

init_attr_dtype = _numpy.dtype([
    ("version", _numpy.int32, ),
    ("mpi_comm", _numpy.intp, ),
    ("args", _numpy.int8, (sizeof(nvshmemx_init_args_t),)),  # opaque
    ], align=True)


cdef class InitAttr:

    cdef:
        readonly object _data

    def __init__(self):
        arr = _numpy.empty(1, dtype=init_attr_dtype)
        self._data = arr.view(_numpy.recarray)
        assert self._data.itemsize == sizeof(nvshmemx_init_attr_t), \
            f"itemsize {self._data.itemsize} mismatches struct size {sizeof(nvshmemx_init_attr_t)}"

    @property
    def ptr(self):
        """Get the pointer address to the data as Python :py:`int`."""
        return self._data.ctypes.data

    @property
    def version(self):
        """version (~_numpy.int32): """
        return int(self._data.version[0])

    @version.setter
    def version(self, val):
        self._data.version = val

    @property
    def mpi_comm(self):
        """mpi_comm (~_numpy.intp): """
        return int(self._data.mpi_comm[0])

    @mpi_comm.setter
    def mpi_comm(self, val):
        self._data.mpi_comm = val


###############################################################################
# Enum
###############################################################################

$enum_defs


###############################################################################
# Error handling
###############################################################################

class NVSHMEMError(Exception):

    def __init__(self, status):
        self.status = status
        cdef str err = f"Status code {status}"
        super(NVSHMEMError, self).__init__(err)

    def __reduce__(self):
        return (type(self), (self.status,))


@cython.profile(False)
cpdef inline check_status(int status):
    if status != 0:
        raise NVSHMEMError(status)


###############################################################################
# Wrapper functions
###############################################################################

$wrapper_defs
