# 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 dataclasses import dataclass

NUMBA_DTYPES = ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64", "float32", "float64", "float16"]

NUMBA_TO_NVSHMEM_TYPES = {
    "int8": "int8",
    "int16": "int16",
    "int32": "int32",
    "int64": "int64",
    "uint8": "uint8",
    "uint16": "uint16",
    "uint32": "uint32",
    "uint64": "uint64",
    "float32": "float",
    "float64": "double",
    "float16": "half",
    "bfloat16": "bfloat16"
}

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
        }

integral_dtypes = [Dtype(numba_dtype_name=dtype, nvshmem_dtype_name=NUMBA_TO_NVSHMEM_TYPES[dtype]) for dtype in NUMBA_DTYPES if dtype in ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]]
float_dtypes = [Dtype(numba_dtype_name=dtype, nvshmem_dtype_name=NUMBA_TO_NVSHMEM_TYPES[dtype]) for dtype in NUMBA_DTYPES if dtype in ["float32", "float64", "float16"]]

rma_dtypes = [Dtype(numba_dtype_name=dtype, nvshmem_dtype_name=NUMBA_TO_NVSHMEM_TYPES[dtype]) for dtype in NUMBA_DTYPES]

@dataclass
class ExecutionScope:
    name: str

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

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


@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,
        }


@dataclass
class ReduceOp:
    name: str

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

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

collectives_execution_scopes = [
    ExecutionScope(name="block"),
    ExecutionScope(name="warp"),
    ExecutionScope(name="thread"),
]

reduction_ops = [
    ReduceOp(name="min"),
    ReduceOp(name="max"),
    ReduceOp(name="sum"),
    ReduceOp(name="prod"),
]

bitwise_reduction_ops =  [
    ReduceOp(name="and"),
    ReduceOp(name="or"),
    ReduceOp(name="xor"),
]

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

@dataclass
class BaseConfig:
    """Parent Class for all APIs
    Defines `to_dict` that converts to Jinja2 template engine friendly dict
    """
    def to_dict(self):
        def to_dict_handler(value):
            if isinstance(value, list):
                return [item.to_dict() for item in value]

        result = {field: to_dict_handler(getattr(self, field)) for field in self.__dataclass_fields__}
        return result

