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

import os
import argparse
from dataclasses import dataclass

import jinja2
from numbast_common import rma_dtypes, Dtype, rma_execution_scopes, ExecutionScope, rma_blocking_semantics, BlockingSemantics, BaseConfig


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


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

    Attributes:
        src_dst: list of Dtypes
        scope: list of ExecutionScopes
        blocking_semantics: list of BlockingSemantics
    """
    src_dst: list[Dtype]
    scope: list[ExecutionScope]
    blocking_semantics: list[BlockingSemantics]


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

    Attributes:
        src_dst: list of Dtypes
        scope: list of ExecutionScopes
        blocking_semantics: list of BlockingSemantics
    """
    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 SignalOpConfig(RMAConfig):
    """Defines all variants of the _signal_op API
    """
    pass


@dataclass
class SignalWaitConfig(RMAConfig):
    """Defines all variants of the _signal_wait API
    """
    pass


@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_signal_op() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _signal_op API"""
    return SignalOpConfig().to_dict()


def generate_signal_wait() -> dict:
    """Populate a jinja2 friendly dict for all variants of the _signal_wait API"""
    return SignalWaitConfig().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.
    
    Config Format:
    {
        "put_config": {
            "src_dst": list[Dtype],
            "scope": list[ExecutionScope],
            "blocking_semantics": list[BlockingSemantics]
        },
        "get_config": {
            "src_dst": list[Dtype],
            "scope": list[ExecutionScope],
            "blocking_semantics": list[BlockingSemantics]
        },
        "put_signal_config": {
            "src_dst": list[Dtype],
            "scope": list[ExecutionScope],
            "blocking_semantics": list[BlockingSemantics]
        },
        "signal_op_config": {
        },
        "signal_wait_config": {
        }
    }

    Returns
    -------
    str: The rendered RMA bindings
    """
    this_dir = os.path.dirname(os.path.abspath(__file__))
    rma_template_path = os.path.join(this_dir, "templates", "core", "device", "numba", "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(),
        "signal_op_config": generate_signal_op(),
        "signal_wait_config": generate_signal_wait(),
        "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 RMA bindings to {target}")
        f.write(generate_rma(config))
