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

import os
import sys
import argparse
from dataclasses import dataclass

import jinja2

this_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(this_dir, "..", "numbast"))

from numbast_common import (  # noqa: E402
    NUMBA_DTYPES, NUMBA_TO_NVSHMEM_TYPES, rma_execution_scopes, ExecutionScope, rma_blocking_semantics,
    BlockingSemantics, BaseConfig,
)

_CUTE_TYPE_NAMES = {
    "int8": "Int8",
    "int16": "Int16",
    "int32": "Int32",
    "int64": "Int64",
    "uint8": "Uint8",
    "uint16": "Uint16",
    "uint32": "Uint32",
    "uint64": "Uint64",
    "float16": "Float16",
    "float32": "Float32",
    "float64": "Float64",
}


class Dtype:

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

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


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


@dataclass
class RMAConfig(BaseConfig):
    """Parent Class for all RMA APIs"""


@dataclass
class PutConfig(RMAConfig):
    """Defines all variants of the _put API"""

    src_dst: list[Dtype]
    scope: list[ExecutionScope]
    blocking_semantics: list[BlockingSemantics]


@dataclass
class GetConfig(RMAConfig):
    """Defines all variants of the _get API"""

    src_dst: list[Dtype]
    scope: list[ExecutionScope]
    blocking_semantics: list[BlockingSemantics]


@dataclass
class PutSignalConfig(RMAConfig):
    """Defines all variants of the _put_signal API"""

    src_dst: list[Dtype]
    scope: list[ExecutionScope]
    blocking_semantics: list[BlockingSemantics]


@dataclass
class PConfig(RMAConfig):
    """Defines all variants of the _p API"""

    src_dst: list[Dtype]


@dataclass
class GConfig(RMAConfig):
    """Defines all variants of the _g API"""

    src_dst: list[Dtype]


def generate_put() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _put API"""
    return PutConfig(
        src_dst=rma_dtypes,
        scope=rma_execution_scopes,
        blocking_semantics=rma_blocking_semantics,
    ).to_dict()


def generate_get() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _get API"""
    return GetConfig(
        src_dst=rma_dtypes,
        scope=rma_execution_scopes,
        blocking_semantics=rma_blocking_semantics,
    ).to_dict()


def generate_put_signal() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _put_signal API"""
    return PutSignalConfig(
        src_dst=rma_dtypes,
        scope=rma_execution_scopes,
        blocking_semantics=rma_blocking_semantics,
    ).to_dict()


def generate_p() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _p API"""
    return PConfig(src_dst=rma_dtypes).to_dict()


def generate_g() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _g API"""
    return GConfig(src_dst=rma_dtypes).to_dict()


def generate_rma(config) -> str:
    """Load the jinja2 template and render it with the config."""
    rma_template_path = os.path.join(this_dir, "templates", "core", "device", "cute", "rma.py.j2")
    with open(rma_template_path, "r") as f:
        template = jinja2.Template(f.read())
    return template.render(**config)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=str, default=".")
    args = parser.parse_args()

    config = {
        "put_config": generate_put(),
        "get_config": generate_get(),
        "put_signal_config": generate_put_signal(),
        "p_config": generate_p(),
        "g_config": generate_g(),
    }
    target = os.path.join(args.output_dir, "rma.py")
    with open(target, "w") as f:
        print(f"Generating CuTe RMA bindings to {target}")
        f.write(generate_rma(config))
