# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

from numba.core import types
from numba.core.extending import overload, overload_method
from numba.core.typing import signature
from numba.cuda import nvvmutils
from numba.cuda.extending import intrinsic
from numba.cuda.types import grid_group, GridGroup as GridGroupClass


class GridGroup:
    """A cooperative group representing the entire grid"""

    def sync() -> None:
        """Synchronize this grid group"""


def this_grid() -> GridGroup:
    """Get the current grid group."""
    return GridGroup()


@intrinsic
def _this_grid(typingctx):
    sig = signature(grid_group)

    def codegen(context, builder, sig, args):
        context.active_code_library.use_cooperative = True
        one = context.get_constant(types.int32, 1)
        mod = builder.module
        return builder.call(
            nvvmutils.declare_cudaCGGetIntrinsicHandle(mod), (one,)
        )

    return sig, codegen


@overload(this_grid, target="cuda")
def _ol_this_grid():
    def impl():
        return _this_grid()

    return impl


@intrinsic
def _grid_group_sync(typingctx, group):
    sig = signature(types.int32, group)

    def codegen(context, builder, sig, args):
        context.active_code_library.use_cooperative = True
        flags = context.get_constant(types.int32, 0)
        mod = builder.module
        return builder.call(
            nvvmutils.declare_cudaCGSynchronize(mod), (*args, flags)
        )

    return sig, codegen


@overload_method(GridGroupClass, "sync", target="cuda")
def _ol_grid_group_sync(group):
    def impl(group):
        return _grid_group_sync(group)

    return impl
