# 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

# Simple wrapper for nvshmemx_team_uniqueid_t typedef
# This is required because Cybind doesn't generate wrappers for classes that are raw typedefs.
cdef class team_uniqueid:
    """Simple wrapper for nvshmemx_team_uniqueid_t (typedef uint64_t)."""
    
    cdef:
        readonly object _data
    
    def __init__(self, value: int = 0):
        """Initialize with an optional value."""
        arr = _numpy.empty(1, dtype=_numpy.uint64)
        self._data = arr.view(_numpy.recarray)
        self._data[0] = value
    
    @property
    def ptr(self):
        """Get the pointer address to the data as Python :py:`int`."""
        return self._data.ctypes.data
    
    @property
    def value(self) -> int:
        """Get the value as Python int."""
        return int(self._data[0])
    
    @value.setter
    def value(self, val: int):
        """Set the value."""
        self._data[0] = val
    
    def __str__(self) -> str:
        return f"team_uniqueid({self.value})"
    
    def __repr__(self) -> str:
        return f"team_uniqueid(value={self.value})"

cdef class TeamUniqueId(team_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


# POD wrapper for nvshmem_team_config_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.

team_config_dtype = _numpy.dtype([
    ("version", _numpy.int32, ),
    ("num_contexts", _numpy.int32, ),
    ("uniqueid", _numpy.uint64, ),
    ("padding", _numpy.int8, (48,)),  # TEAM_CONFIG_V2_PADDING = 48
    ], align=True)


cdef class TeamConfig:

    cdef:
        readonly object _data

    def __init__(self):
        arr = _numpy.empty(1, dtype=team_config_dtype)
        self._data = arr.view(_numpy.recarray)
        assert self._data.itemsize == sizeof(nvshmem_team_config_t), \
            f"itemsize {self._data.itemsize} mismatches struct size {sizeof(nvshmem_team_config_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 num_contexts(self):
        """num_contexts (~_numpy.int32): """
        return int(self._data.num_contexts[0])

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

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

    @uniqueid.setter
    def uniqueid(self, val):
        self._data.uniqueid = 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
