# 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

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))