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"),
]