# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from dataclasses import dataclass

from numba_constants import NUMBA_DTYPES, NUMBA_TO_NVSHMEM_TYPES, NUMBA_TO_NVSHMEM_TYPES_AMO


class Dtype:

    def __init__(self, numba_dtype_name, nvshmem_dtype_name):
        self.numba_dtype_name = numba_dtype_name
        self.nvshmem_dtype_name = nvshmem_dtype_name

    def to_dict(self):
        return {"numba_type_name": self.numba_dtype_name, "nvshmem_type_name": self.nvshmem_dtype_name}


def get_dtype(type_name: str, mapping: dict[str, str] = NUMBA_TO_NVSHMEM_TYPES_AMO) -> Dtype:
    return Dtype(numba_dtype_name=type_name, nvshmem_dtype_name=mapping[type_name])


rma_dtypes = [Dtype(numba_dtype_name=dtype, nvshmem_dtype_name=NUMBA_TO_NVSHMEM_TYPES[dtype]) for dtype in NUMBA_DTYPES]
amo_std_dtypes = [get_dtype("uint64"), get_dtype("uint32"), get_dtype("int32"), get_dtype("int64")]
amo_ext_dtypes = [get_dtype("float32"), get_dtype("float64")] + amo_std_dtypes
amo_bit_dtypes = amo_std_dtypes


@dataclass
class ExecutionScope:
    name: str

    @property
    def suffix(self):
        if self.name == "device":
            return ""
        else:
            return f"_{self.name}"

    def to_dict(self):
        return {
            "suffix": self.suffix,
        }


rma_execution_scopes = [
    ExecutionScope(name="device"),
    ExecutionScope(name="block"),
    ExecutionScope(name="warp"),
]


@dataclass
class BlockingSemantics:
    kind: str

    @property
    def name(self):
        if self.kind == "blocking":
            return ""
        else:
            return f"_nbi"

    def to_dict(self):
        return {
            "name": self.name,
        }


rma_blocking_semantics = [
    BlockingSemantics(kind="blocking"),
    BlockingSemantics(kind="nonblocking"),
]
