# 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

from nvshmem.bindings.device.numba import *

import cffi
from numba.core.extending import overload
from numba.types import int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, float16, Array
from numba.cuda.types import bfloat16

__all__ = ["p", "g", "put", "get", "put_nbi", "get_nbi", "put_block", "get_block", "put_nbi_block", "get_nbi_block", "put_warp", "get_warp", "put_nbi_warp", "get_nbi_warp", "put_signal_block", "put_signal", "put_signal_nbi", "put_signal_warp", "put_signal_nbi_block", "put_signal_nbi_warp", "put_signal_nbi_warp"]

# TODO: create a global ffi object for other high level bindings to use
ffi = cffi.FFI()

def put(src: Array, dst: Array, pe: int32) -> None:
    """
    Copies data from local ``src`` to symmetric ``dst`` on PE ``pe``. 
    This is a thread-level operation.

    Args:
        - ``src`` (``Array``): Local source array on this PE to copy from.
        - ``dst`` (``Array``): Symmetric destination array on PE ``pe`` to copy to.
        - ``pe`` (``int``): PE to copy to.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def get(dst: Array, src: Array, pe: int32) -> None:
    """
    Copies data from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a thread-level operation.

    Args:
        - ``dst`` (``Array``): Local destination array on this PE to copy to.
        - ``src`` (``Array``): Symmetric source array from PE ``pe`` to copy from.
        - ``pe`` (``int``): PE to copy from.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def put_nbi(src: Array, dst: Array, pe: int32) -> None:
    """
    Non-blockingly copies data from local ``src`` to symmetric ``dst`` on PE ``pe``.
    This is a thread-level operation.

    Args:
        - ``src`` (``Array``): Local source array on this PE to copy from.
        - ``dst`` (``Array``): Symmetric destination array on PE ``pe`` to copy to.
        - ``pe`` (``int``): PE to copy to.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def get_nbi(dst: Array, src: Array, pe: int32) -> None:
    """
    Non-blockingly copies data from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a thread-level operation.

    Args:
        - ``dst`` (``Array``): Local destination array on this PE to copy to.
        - ``src`` (``Array``): Symmetric source array from PE ``pe`` to copy from.
        - ``pe`` (``int``): PE to copy from.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def put_block(src: Array, dst: Array, pe: int32) -> None:
    """
    Copies from local ``src`` to symmetric ``dst`` on PE ``pe``.
    This is a CTA-level operation. All threads in the CTA must call this function with the same arguments.

    Args:
        - ``src`` (``Array``): Local source array on this PE to copy from.
        - ``dst`` (``Array``): Symmetric destination array on PE ``pe`` to copy to.
        - ``pe`` (``int``): PE to copy to.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def get_block(dst: Array, src: Array, pe: int32) -> None:
    """
    Copies from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a CTA-level operation. All threads in the CTA must call this function with the same arguments.

    Args:
        - ``dst`` (``Array``): Local destination array on this PE to copy to.
        - ``src`` (``Array``): Symmetric source array from PE ``pe`` to copy from.
        - ``pe`` (``int``): PE to copy from.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def put_nbi_block(src: Array, dst: Array, pe: int32) -> None:
    """
    Non-blockingly copies from local ``src`` to symmetric ``dst`` on PE ``pe``.
    This is a CTA-level operation. All threads in the CTA must call this function with the same arguments.

    Args:
        - ``src`` (``Array``): Local source array on this PE to copy from.
        - ``dst`` (``Array``): Symmetric destination array on PE ``pe`` to copy to.
        - ``pe`` (``int``): PE to copy to.
    """
    pass

def get_nbi_block(dst: Array, src: Array, pe: int32) -> None:
    """
    Non-blockingly copies from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a CTA-level operation. All threads in the CTA must call this function with the same arguments.

    Args:
        - ``dst`` (``Array``): Local destination array on this PE to copy to.
        - ``src`` (``Array``): Symmetric source array from PE ``pe`` to copy from.
        - ``pe`` (``int``): PE to copy from.
    """
    pass

def put_warp(src: Array, dst: Array, pe: int32) -> None:
    """
    Copies from local ``src`` to symmetric ``dst`` on PE ``pe``.
    This is a warp-level operation. All threads in the warp must call this function with the same arguments.

    Args:
        - ``src`` (``Array``): Local source array on this PE to copy from.
        - ``dst`` (``Array``): Symmetric destination array on PE ``pe`` to copy to.
        - ``pe`` (``int``): PE to copy to.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def get_warp(dst: Array, src: Array, pe: int32) -> None:
    """
    Copies from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a warp-level operation. All threads in the warp must call this function with the same arguments.

    Args:
        - ``dst`` (``Array``): Local destination array on this PE to copy to.
        - ``src`` (``Array``): Symmetric source array from PE ``pe`` to copy from.
        - ``pe`` (``int``): PE to copy from.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def put_nbi_warp(src: Array, dst: Array, pe: int32) -> None:
    """
    Non-blockingly copies from local ``src`` to symmetric ``dst`` on PE ``pe``.
    This is a warp-level operation. All threads in the warp must call this function with the same arguments.

    Args:
        - ``src`` (``Array``): Local source array on this PE to copy from.
        - ``dst`` (``Array``): Symmetric destination array on PE ``pe`` to copy to.
        - ``pe`` (``int``): PE to copy to.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

def get_nbi_warp(dst: Array, src: Array, pe: int32) -> None:
    """
    Non-blockingly copies from symmetric ``src`` on PE ``pe`` to local ``dst``.
    This is a warp-level operation. All threads in the warp must call this function with the same arguments.

    Args:
        - ``dst`` (``Array``): Local destination array on this PE to copy to.
        - ``src`` (``Array``): Symmetric source array from PE ``pe`` to copy from.
        - ``pe`` (``int``): PE to copy from.
    
    Note:
        When ``src`` and ``dst`` are of different sizes, only data of the smaller size is copied.
        ``src`` and ``dst`` must have the same data type.
    """
    pass

# put variations
{%- for scope in put_config.scope %}
{%- for blocking_semantics in put_config.blocking_semantics %}
@overload(put{{ blocking_semantics.name }}{{ scope.suffix }})
def put{{ blocking_semantics.name }}{{ scope.suffix }}_ol(dst, src, pe):
{%- for dtype in put_config.src_dst %}
    {{ "if" if loop.first else "elif" }} dst == Array(dtype={{dtype.numba_type_name}}, ndim=dst.ndim, layout=dst.layout) and src == Array(dtype={{dtype.numba_type_name}}, ndim=src.ndim, layout=src.layout):
        def impl(dst, src, pe):
            dstptr = ffi.from_buffer(dst)
            srcptr = ffi.from_buffer(src)
            nelems = min(dst.size, src.size)
            {{dtype.nvshmem_type_name}}_put{{blocking_semantics.name}}{{scope.suffix}}(dstptr, srcptr, nelems, pe)
        return impl
{%- endfor %}
{% endfor %}
{% endfor %}

# get variations
{%- for scope in get_config.scope %}
{%- for blocking_semantics in get_config.blocking_semantics %}
@overload(get{{ blocking_semantics.name }}{{ scope.suffix }})
def get{{ blocking_semantics.name }}{{ scope.suffix }}_ol(dst, src, pe):
{%- for dtype in get_config.src_dst %}
    {{ "if" if loop.first else "elif" }} dst == Array(dtype={{dtype.numba_type_name}}, ndim=dst.ndim, layout=dst.layout) and src == Array(dtype={{dtype.numba_type_name}}, ndim=src.ndim, layout=src.layout):
        def impl(dst, src, pe):
            dstptr = ffi.from_buffer(dst)
            srcptr = ffi.from_buffer(src)
            nelems = min(dst.size, src.size)
            {{dtype.nvshmem_type_name}}_get{{blocking_semantics.name}}{{scope.suffix}}(dstptr, srcptr, nelems, pe)
        return impl
{%- endfor %}
{% endfor %}
{% endfor %}

# put_signal variations
{%- for scope in put_signal_config.scope %}
{%- for blocking_semantics in put_signal_config.blocking_semantics %}
def put_signal{{ blocking_semantics.name }}{{ scope.suffix }}():
    """
    Put data with a signal operation from ``src`` to ``dst`` on PE ``pe`` in a {{blocking_semantics.name}} manner at a {{scope.name}} scope. Device initiated.

    Signal variables must be an ``Array`` of dtype ``int64`` (8 bytes) allocated by or registered with NVSHMEM4Py.
    Supported signal operations are ``SignalOp.SIGNAL_SET`` and ``SignalOp.SIGNAL_ADD``.

    Args:
        - ``dst`` (``Array``): Destination symmetric array on remote PE.
        - ``src`` (``Array``): Source symmetric array on this PE.
        - ``signal_var`` (``Array``): Symmetric signal variable (dtype ``int64``).
        - ``signal_val`` (``int``): Signal value.
        - ``signal_op`` (``SignalOp``): Signal operation type.
        - ``pe`` (``int``): Target PE to put to.
    """
    pass
@overload(put_signal{{ blocking_semantics.name }}{{ scope.suffix }})
def put_signal{{ blocking_semantics.name }}{{ scope.suffix }}_ol(dst, src, signal_var, signal_val, signal_op, pe):
{%- for dtype in put_signal_config.src_dst %}
    {{ "if" if loop.first else "elif" }} dst == Array(dtype={{dtype.numba_type_name}}, ndim=dst.ndim, layout=dst.layout) and src == Array(dtype={{dtype.numba_type_name}}, ndim=src.ndim, layout=src.layout):
        def impl(dst, src, signal_var, signal_val, signal_op, pe):
            dstptr = ffi.from_buffer(dst)
            srcptr = ffi.from_buffer(src)
            signal_varptr = ffi.from_buffer(signal_var)
            nelems = min(dst.size, src.size)
            {{dtype.nvshmem_type_name}}_put_signal{{blocking_semantics.name}}{{scope.suffix}}(dstptr, srcptr, nelems, signal_varptr, signal_val, signal_op, pe)
        return impl
{%- endfor %}
{% endfor %}
{% endfor %}

# p variations
def p():
    """
    Put immediate data from src to dst on PE pe. Device initiated.

    ``src`` must be a scalar value, passed as a symmetric Array of size 1.

    Args:
        - ``dst`` (``Array``): Destination symmetric array on remote PE.
        - ``src`` (``Array``): Scalar source value (size-1 array) on this PE.
        - ``pe`` (``int``): PE to put to.
    """
    pass
@overload(p)
def p_ol(dst, src, pe):
{%- for dtype in p_config.src_dst %}
    {{ "if" if loop.first else "elif" }} dst == Array(dtype={{dtype.numba_type_name}}, ndim=dst.ndim, layout=dst.layout):
        def impl(dst, src, pe):
            dstptr = ffi.from_buffer(dst)
            {{dtype.nvshmem_type_name}}_p(dstptr, src, pe)
        return impl
{%- endfor %}

def g():
    """
    Get immediate data from src on PE pe to local dst. Device initiated.

    ``src`` must be a scalar value, passed as a symmetric Array of size 1.

    Args:
        - ``src`` (``Array``): Source symmetric array (size-1) on remote PE.
        - ``pe`` (``int``): PE to get from.
    Returns:
        Scalar value retrieved from ``src`` on PE ``pe``.
    """
    pass
@overload(g)
def g_ol(src, pe):
{%- for dtype in p_config.src_dst %}
    {{ "if" if loop.first else "elif" }} src == Array(dtype={{dtype.numba_type_name}}, ndim=src.ndim, layout=src.layout):
        def impl(src, pe):
            srcptr = ffi.from_buffer(src)
            return {{dtype.nvshmem_type_name}}_g(srcptr, pe)
        return impl
{%- endfor %}
