# 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

import nvshmem.core
import nvshmem.bindings.device.numba as bindings

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


__all__ = [ "sync_block", "sync_warp", "sync", "sync_all", "sync_all_block", "sync_all_warp", 
            "barrier", "barrier_block", "barrier_warp", "barrier_all", "barrier_all_block", "barrier_all_warp", 
            "reduce", "reduce_block", "reduce_warp",
            "reducescatter", "reducescatter_block", "reducescatter_warp", 
            "fcollect", "fcollect_block", "fcollect_warp", 
            "broadcast", "broadcast_block", "broadcast_warp",
            "alltoall", "alltoall_block", "alltoall_warp"]

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

def sync(team: nvshmem.core.Teams):
    """
    Executes a sync across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
    """
    pass

def sync_block(team: nvshmem.core.Teams):
    """
    Executes a block-wide sync across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
    """
    pass

def sync_warp(team: nvshmem.core.Teams):
    """
    Executes a warp-wide sync across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
    """
    pass

def sync_all():
    """
    Executes a thread-wide sync across all PEs in the runtime.
    """
    pass

def sync_all_block():
    """
    Executes a block-wide sync across all PEs in the runtime.
    """
    pass

def sync_all_warp():
    """
    Executes a warp-wide sync across all PEs in the runtime.
    """
    pass

def barrier(team: nvshmem.core.Teams):
    """
    Executes a thread-wide barrier across the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
    """
    pass

def barrier_block(team: nvshmem.core.Teams):
    """
    Executes a block-wide barrier across all threads in the block.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
    """
    pass

def barrier_warp(team: nvshmem.core.Teams):
    """
    Executes a warp-wide barrier across all threads in the warp.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
    """
    pass

def barrier_all():
    """
    Executes a barrier across all PEs in the runtime.
    """
    pass


def barrier_all_block():
    """
    Executes a block-wide barrier across all PEs in the runtime.
    """
    pass

def barrier_all_warp():
    """
    Executes a warp-wide barrier across all PEs in the runtime.
    """
    pass

def reduce(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, op: str):
    """
    Performs a reduction from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the reduction routine.
        - ``src_array`` (``Array``): Symmetric array that contains at least one element for each separate reduction routine.
        - ``op`` (``str``): String representing the reduction operator.
    
    Supported reduction operators:
    See https://docs.nvidia.com/nvshmem/api/gen/api/collectives.html?highlight=allreduce#nvshmem-reductions
    for supported reduction operators.

    Note:
    In case that ``src_array`` and ``dst_array`` are of different sizes, the size of the smaller array is reduced.
    """
    pass

def reduce_block(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, op: str):
    """
    Performs a block-wide reduction from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the reduction routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
        - ``op`` (``str``): String representing the reduction operator.
    
    Supported reduction operators:
    See https://docs.nvidia.com/nvshmem/api/gen/api/collectives.html?highlight=allreduce#nvshmem-reductions
    for supported reduction operators.

    Note:
    In case that ``src_array`` and ``dst_array`` are of different sizes, the size of the smaller array is reduced.
    """
    pass

def reduce_warp(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, op: str):
    """
    Performs a warp-wide reduction from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the reduction routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
        - ``op`` (``str``): String representing the reduction operator.
    
    Supported reduction operators:
    See https://docs.nvidia.com/nvshmem/api/gen/api/collectives.html?highlight=allreduce#nvshmem-reductions
    for supported reduction operators.

    Note:
    In case that ``src_array`` and ``dst_array`` are of different sizes, the size of the smaller array is reduced.
    """
    pass

def reducescatter(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, op: str):
    """
    Performs a thread-scoped reduce-scatter operation from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the reduce-scatter routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
        - ``op`` (``str``): String representing the reduction operator.
    
    Supported reduction operators:
    See https://docs.nvidia.com/nvshmem/api/gen/api/collectives.html?highlight=allreduce#nvshmem-reductions
    for supported reduction operators.

    Note:
    Array size is taken from the ``dst_array``. The ``dst_array`` must be >= ``(src_array.size // team_n_pes(team))``.
    """
    pass

def reducescatter_block(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, op: str):
    """
    Performs a block-scoped reduce-scatter operation from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the reduce-scatter routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
        - ``op`` (``str``): String representing the reduction operator.
    
    Supported reduction operators:
    See https://docs.nvidia.com/nvshmem/api/gen/api/collectives.html?highlight=allreduce#nvshmem-reductions
    for supported reduction operators.

    Note:
    Array size is taken from the ``dst_array``. The ``src_array`` must be >= ``(dst_array.size * team_n_pes(team))``.
    """
    pass

def reducescatter_warp(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, op: str):
    """
    Performs a warp-scoped reduce-scatter operation from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the reduce-scatter routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
        - ``op`` (``str``): String representing the reduction operator.
    
    Supported reduction operators:
    See https://docs.nvidia.com/nvshmem/api/gen/api/collectives.html?highlight=allreduce#nvshmem-reductions
    for supported reduction operators.

    Note:
    Array size is taken from the ``dst_array``. The ``src_array`` must be >= ``(src_array.size // team_n_pes(team))``.
    """
    pass

def fcollect(team: nvshmem.core.Teams, dst_array: Array, src_array: Array):
    """
    Performs a thread-scoped fcollect operation from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the fcollect routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
    
    Note:
    Array size is taken from the ``src_array``. The ``src_array`` must be >= ``(src_array.size * team_n_pes(team))``.
    """
    pass

def fcollect_block(team: nvshmem.core.Teams, dst_array: Array, src_array: Array):
    """
    Performs a block-scoped fcollect operation from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the fcollect routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
    
    Note:
    Array size is taken from the ``src_array``. The ``dst_array`` must be >= ``(src_array.size * team_n_pes(team))``.
    """
    pass
    
def fcollect_warp(team: nvshmem.core.Teams, dst_array: Array, src_array: Array):
    """
    Performs a warp-scoped fcollect operation from src_array to dst_array across all PEs in the team.

    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Symmetric array to store the result of the fcollect routine.
        - ``src_array`` (``Array``): Symmetric array that contains the source data.
    
    Note:
    Array size is taken from the ``src_array``. The ``dst_array`` must be >= ``(src_array.size * team_n_pes(team))``.
    """
    pass
    
def broadcast(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, root: int=0):
    """
    Performs a thread-scoped broadcast operation from src_array to dst_array across all PEs in the team.
    
    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Destination symmetric array.
        - ``src_array`` (``Array``): Source symmetric array.
        - ``root`` (``int``): Root PE for the broadcast.
    """
    pass
    
def broadcast_block(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, root: int=0):
    """
    Performs a block-scoped broadcast operation from src_array to dst_array across all PEs in the team.
    
    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Destination symmetric array.
        - ``src_array`` (``Array``): Source symmetric array.
        - ``root`` (``int``): Root PE for the broadcast.
    """
    pass
    
def broadcast_warp(team: nvshmem.core.Teams, dst_array: Array, src_array: Array, root: int=0):
    """
    Performs a warp-scoped broadcast operation from src_array to dst_array across all PEs in the team.
    
    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Destination symmetric array.
        - ``src_array`` (``Array``): Source symmetric array.
        - ``root`` (``int``): Root PE for the broadcast.
    """
    pass
    
def alltoall(team: nvshmem.core.Teams, dst_array: Array, src_array: Array):
    """
    Performs a thread-scoped alltoall operation from src_array to dst_array across all PEs in the team.
    
    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Destination symmetric array.
        - ``src_array`` (``Array``): Source symmetric array.
    """
    pass
    
def alltoall_block(team: nvshmem.core.Teams, dst_array: Array, src_array: Array):
    """
    Performs a block-scoped alltoall operation from src_array to dst_array across all PEs in the team.
    
    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Destination symmetric array.
        - ``src_array`` (``Array``): Source symmetric array.
    """
    pass

def alltoall_warp(team: nvshmem.core.Teams, dst_array: Array, src_array: Array):
    """
    Performs a warp-scoped alltoall operation from src_array to dst_array across all PEs in the team.
    
    Args:
        - ``team`` (``Teams``): NVSHMEM team handle.
        - ``dst_array`` (``Array``): Destination symmetric array.
        - ``src_array`` (``Array``): Source symmetric array.
    """
    pass

# sync variations
{%- for scope in sync_config.scope %}
@overload(sync{{ scope.suffix }})
def sync{{ scope.suffix }}_ol(team):
    def impl(team):
        bindings.team_sync{{scope.suffix}}(team)
    return impl
{%- endfor %}

# sync_all variations
{%- for scope in sync_all_config.scope %}
@overload(sync_all{{ scope.suffix }})
def sync_all{{ scope.suffix }}_ol():
    def impl():
        bindings.sync_all{{scope.suffix}}()
    return impl
{%- endfor %}

# barrier variations
{%- for scope in barrier_config.scope %}
@overload(barrier{{ scope.suffix }})
def barrier{{ scope.suffix }}_ol(team):
    def impl(team):
        bindings.barrier{{scope.suffix}}(team)
    return impl
{%- endfor %}

# barrier_all variations
{%- for scope in barrier_all_config.scope %}
@overload(barrier_all{{ scope.suffix }})
def barrier_all{{ scope.suffix }}_ol():
    def impl():
        bindings.barrier_all{{scope.suffix}}()
    return impl
{%- endfor %}

# reduce variations
{%- for scope in reduce_config.scope %}
@overload(reduce{{ scope.suffix }})
def reduce{{ scope.suffix }}_ol(team, dst_array, src_array, op):
    if isinstance(op, types.Literal) and isinstance(op.literal_value, str):
        if isinstance(dst_array.dtype, types.Integer) and op.literal_value not in { {{ '"' ~ (((reduce_config.bitwise_op_literal | map(attribute='name') | list) + (reduce_config.reduction_op_literal | map(attribute='name') | list)) | join('", "')) ~ '"' }} }:
            raise ValueError(f"Invalid reduction operator: {op.literal_value} for integer data type")
        if isinstance(dst_array.dtype, types.Float) and op.literal_value not in { {{ '"' ~ (reduce_config.reduction_op_literal | map(attribute='name') | list | join('", "')) ~ '"' }} }:
            raise ValueError(f"Invalid reduction operator: {op.literal_value} for float data type")

        {%- for op_literal in reduce_config.reduction_op_literal %}
        {{ "if" if loop.first else "elif" }} op.literal_value == "{{ op_literal.name }}":
            {%- for dtype in reduce_config.all_dtypes %}
            {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
                def impl(team, dst_array, src_array, op):
                    dst_ptr = ffi.from_buffer(dst_array)
                    src_ptr = ffi.from_buffer(src_array)
                    nelem = dst_array.size
                    bindings.{{dtype.nvshmem_type_name}}_{{op_literal.name}}_reduce{{scope.suffix}}(team, dst_ptr, src_ptr, nelem)
                return impl
            {%- endfor %}
        {%- endfor %}

        {%- for op_literal in reduce_config.bitwise_op_literal %}
        {{ "if" if loop.first else "elif" }} op.literal_value == "{{ op_literal.name }}":
            {%- for dtype in reduce_config.integral_dtypes %}
            {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
                def impl(team, dst_array, src_array, op):
                    dst_ptr = ffi.from_buffer(dst_array)
                    src_ptr = ffi.from_buffer(src_array)
                    nelem = dst_array.size
                    bindings.{{dtype.nvshmem_type_name}}_{{op_literal.name}}_reduce{{scope.suffix}}(team, dst_ptr, src_ptr, nelem)
                return impl
            {%- endfor %}
        {%- endfor %}

{%- endfor %}

# reducescatter variations
{%- for scope in reducescatter_config.scope %}
@overload(reducescatter{{ scope.suffix }})
def reducescatter{{ scope.suffix }}_ol(team, dst_array, src_array, op):
    if isinstance(op, types.Literal) and isinstance(op.literal_value, str):
        if isinstance(dst_array.dtype, types.Integer) and op.literal_value not in {"and", "or", "xor", "min", "max", "sum", "prod"}:
            raise ValueError(f"Invalid reduction operator: {op.literal_value} for integer data type")
        if isinstance(dst_array.dtype, types.Float) and op.literal_value not in {"min", "max", "sum", "prod"}:
            raise ValueError(f"Invalid reduction operator: {op.literal_value} for float data type")

        {%- for op_literal in reducescatter_config.reduction_op_literal %}
        {{ "if" if loop.first else "elif" }} op.literal_value == "{{ op_literal.name }}":
            {%- for dtype in reduce_config.all_dtypes %}
            {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
                def impl(team, dst_array, src_array, op):
                    dst_ptr = ffi.from_buffer(dst_array)
                    src_ptr = ffi.from_buffer(src_array)
                    nelem = dst_array.size
                    bindings.{{dtype.nvshmem_type_name}}_{{op_literal.name}}_reducescatter{{scope.suffix}}(team, dst_ptr, src_ptr, nelem)
                return impl
            {%- endfor %}
        {%- endfor %}

        {%- for op_literal in reducescatter_config.bitwise_op_literal %}
        {{ "if" if loop.first else "elif" }} op.literal_value == "{{ op_literal.name }}":
            {%- for dtype in reducescatter_config.integral_dtypes %}
            {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
                def impl(team, dst_array, src_array, op):
                    dst_ptr = ffi.from_buffer(dst_array)
                    src_ptr = ffi.from_buffer(src_array)
                    nelem = dst_array.size
                    bindings.{{dtype.nvshmem_type_name}}_{{op_literal.name}}_reducescatter{{scope.suffix}}(team, dst_ptr, src_ptr, nelem)
                return impl
            {%- endfor %}
        {%- endfor %}

{%- endfor %}

# fcollect variations
{%- for scope in fcollect_config.scope %}
@overload(fcollect{{ scope.suffix }})
def fcollect{{ scope.suffix }}_ol(team, dst_array, src_array):
    {%- for dtype in fcollect_config.all_dtypes %}
    {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
        def impl(team, dst_array, src_array):
            dst_ptr = ffi.from_buffer(dst_array)
            src_ptr = ffi.from_buffer(src_array)
            nelem = src_array.size
            bindings.{{dtype.nvshmem_type_name}}_fcollect{{scope.suffix}}(team, dst_ptr, src_ptr, nelem)
        return impl
    {%- endfor %}

{%- endfor %}

# broadcast variations
{%- for scope in broadcast_config.scope %}
@overload(broadcast{{ scope.suffix }})
def broadcast{{ scope.suffix }}_ol(team, dst_array, src_array, root=0):
    {%- for dtype in broadcast_config.all_dtypes %}
    {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
        def impl(team, dst_array, src_array, root=0):
            dst_ptr = ffi.from_buffer(dst_array)
            src_ptr = ffi.from_buffer(src_array)
            nelem = min(dst_array.size, src_array.size)
            bindings.{{dtype.nvshmem_type_name}}_broadcast{{scope.suffix}}(team, dst_ptr, src_ptr, nelem, root)
        return impl
    {%- endfor %}

{%- endfor %}

# alltoall variations
{%- for scope in alltoall_config.scope %}
@overload(alltoall{{ scope.suffix }})
def alltoall{{ scope.suffix }}_ol(team, dst_array, src_array):
    {%- for dtype in alltoall_config.all_dtypes %}
    {{ "if" if loop.first else "elif" }} dst_array == Array(dtype={{dtype.numba_type_name}}, ndim=dst_array.ndim, layout=dst_array.layout) and src_array == Array(dtype={{dtype.numba_type_name}}, ndim=src_array.ndim, layout=src_array.layout):
        def impl(team, dst_array, src_array):
            dst_ptr = ffi.from_buffer(dst_array)
            src_ptr = ffi.from_buffer(src_array)
            # alltoall is a special case because it takes the per-device count.
            nelem = src_array.size // bindings.team_n_pes(team)
            bindings.{{dtype.nvshmem_type_name}}_alltoall{{scope.suffix}}(team, dst_ptr, src_ptr, nelem)
        return impl
    {%- endfor %}

{%- endfor %}
