# 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 parameters import amo_std_dtypes, amo_ext_dtypes, amo_bit_dtypes, Dtype, rma_execution_scopes, ExecutionScope, rma_blocking_semantics

# This is the parent class of all AMO configs, which generalizes the common fields
@dataclass
class AmoConfig:
    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

@dataclass
class AtomicFetchConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicSetConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicCompareSwapConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicFetchIncConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicIncConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicFetchAddConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicAddConfig(AmoConfig):
    src_dst: list[Dtype]


@dataclass
class AtomicAndConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicSwapConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicFetchAndConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicFetchOrConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicOrConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicFetchXorConfig(AmoConfig):
    src_dst: list[Dtype]

@dataclass
class AtomicXorConfig(AmoConfig):
    src_dst: list[Dtype]

def generate_atomic_fetch() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AMO APIs"""
    return AtomicFetchConfig(src_dst=amo_ext_dtypes).to_dict()

def generate_atomic_set() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicSet APIs"""
    return AtomicSetConfig(src_dst=amo_ext_dtypes).to_dict()

def generate_atomic_fetch_inc() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicFetchInc APIs"""
    return AtomicFetchIncConfig(src_dst=amo_std_dtypes).to_dict()

def generate_atomic_inc() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicInc APIs"""
    return AtomicIncConfig(src_dst=amo_std_dtypes).to_dict()

def generate_atomic_fetch_add() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicFetchAdd APIs"""
    return AtomicFetchAddConfig(src_dst=amo_std_dtypes).to_dict()

def generate_atomic_add() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicAdd APIs"""
    return AtomicAddConfig(src_dst=amo_std_dtypes).to_dict()

def generate_atomic_and() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicAnd APIs"""
    return AtomicAndConfig(src_dst=amo_bit_dtypes).to_dict()

def generate_atomic_swap() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicSwap APIs"""
    return AtomicSwapConfig(src_dst=amo_ext_dtypes).to_dict()

def generate_atomic_fetch_and() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicFetchAnd APIs"""
    return AtomicFetchAndConfig(src_dst=amo_bit_dtypes).to_dict()

def generate_atomic_or() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicOr APIs"""
    return AtomicOrConfig(src_dst=amo_bit_dtypes).to_dict()

def generate_atomic_fetch_or() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicFetchOr APIs"""
    return AtomicFetchOrConfig(src_dst=amo_bit_dtypes).to_dict()

def generate_atomic_xor() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicXor APIs"""
    return AtomicXorConfig(src_dst=amo_bit_dtypes).to_dict()

def generate_atomic_fetch_xor() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicFetchXor APIs"""
    return AtomicFetchXorConfig(src_dst=amo_bit_dtypes).to_dict()

def generate_atomic_compare_swap() -> dict:
    """Populate a jinja2 friendly dict for all variants of the AtomicCompareSwap APIs"""
    return AtomicCompareSwapConfig(src_dst=amo_std_dtypes).to_dict()

def generate_amo(config) -> str:
    """Load the jinja2 template and render it with the config.
    
    Config Format:
    {
        "atomic_fetch_config": {
            "src_dst": list[Dtype]
        },
    }
    """
    this_dir = os.path.dirname(os.path.abspath(__file__))
    amo_template_path = os.path.join(this_dir, "templates", "core", "device", "numba", "amo.py.j2")
    with open(amo_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 = {
        "atomic_fetch_config": generate_atomic_fetch(),
        "atomic_set_config": generate_atomic_set(),
        "atomic_compare_swap_config": generate_atomic_compare_swap(),
        "atomic_fetch_inc_config": generate_atomic_fetch_inc(),
        "atomic_inc_config": generate_atomic_inc(),
        "atomic_fetch_add_config": generate_atomic_fetch_add(),
        "atomic_add_config": generate_atomic_add(),
        "atomic_and_config": generate_atomic_and(),
        "atomic_swap_config": generate_atomic_swap(),
        "atomic_fetch_and_config": generate_atomic_fetch_and(),
        "atomic_or_config": generate_atomic_or(),
        "atomic_fetch_or_config": generate_atomic_fetch_or(),
        "atomic_xor_config": generate_atomic_xor(),
        "atomic_fetch_xor_config": generate_atomic_fetch_xor(),
    }
    target = os.path.join(args.output_dir, "amo.py")
    with open(target, "w") as f:
        print(f"Generating RMA bindings to {target}")
        f.write(generate_amo(config))