# 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 collectives_execution_scopes, integral_dtypes, float_dtypes, reduction_ops, bitwise_reduction_ops, ExecutionScope, BaseConfig, Dtype, ReduceOp

@dataclass
class CollectivesConfig(BaseConfig):
    """Parent Class for all Collectives APIs"""
    pass

@dataclass
class SyncConfig(CollectivesConfig):
    """Defines all variants of the sync_ API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]

@dataclass
class SyncAllConfig(CollectivesConfig):
    """Defines all variants of the sync_all API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]

@dataclass
class BarrierConfig(CollectivesConfig):
    """Defines all variants of the barrier API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]

@dataclass
class BarrierAllConfig(CollectivesConfig):
    """Defines all variants of the barrier_all API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]

@dataclass
class ReduceConfig(CollectivesConfig):
    """Defines all variants of the reduce API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]
    integral_dtypes: list[Dtype]
    reduction_op_literal: list[ReduceOp]
    bitwise_op_literal: list[ReduceOp]

@dataclass
class ReducescatterConfig(CollectivesConfig):
    """Defines all variants of the reduce API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]
    integral_dtypes: list[Dtype]
    reduction_op_literal: list[ReduceOp]
    bitwise_op_literal: list[ReduceOp]

@dataclass
class FCollectConfig(CollectivesConfig):
    """Defines all variants of the fcollect API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]

@dataclass
class BroadcastConfig(CollectivesConfig):
    """Defines all variants of the broadcast API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]

@dataclass
class AlltoallConfig(CollectivesConfig):
    """Defines all variants of the alltoall API

    Attributes:
        scope: list of ExecutionScopes
    """
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]

def generate_broadcast() -> dict:
    """Populate a jinja2 friendly dict for all variants of the broadcast API"""
    return BroadcastConfig(scope=collectives_execution_scopes, all_dtypes=integral_dtypes + float_dtypes).to_dict()

def generate_alltoall() -> dict:
    """Populate a jinja2 friendly dict for all variants of the alltoall API"""
    return AlltoallConfig(scope=collectives_execution_scopes, all_dtypes=integral_dtypes + float_dtypes).to_dict()

def generate_reducescatter() -> dict:
    """Populate a jinja2 friendly dict for all variants of the reduce API"""
    return ReducescatterConfig(scope=collectives_execution_scopes, all_dtypes=integral_dtypes + float_dtypes, integral_dtypes=integral_dtypes, reduction_op_literal=reduction_ops, bitwise_op_literal=bitwise_reduction_ops).to_dict()

def generate_fcollect() -> dict:
    """Populate a jinja2 friendly dict for all variants of the fcollect API"""
    return FCollectConfig(scope=collectives_execution_scopes, all_dtypes=integral_dtypes + float_dtypes).to_dict()

def generate_sync() -> dict:
    """Populate a jinja2 friendly dict for all variants of the sync_ API"""
    return SyncConfig(scope=collectives_execution_scopes).to_dict()


def generate_sync_all() -> dict:
    """Populate a jinja2 friendly dict for all variants of the sync_all API"""
    return SyncAllConfig(scope=collectives_execution_scopes).to_dict()


def generate_barrier() -> dict:
    """Populate a jinja2 friendly dict for all variants of the barrier API"""
    return BarrierConfig(scope=collectives_execution_scopes).to_dict()

def generate_barrier_all() -> dict:
    """Populate a jinja2 friendly dict for all variants of the barrier_all API"""
    return BarrierAllConfig(scope=collectives_execution_scopes).to_dict()

def generate_reduce() -> dict:
    """Populate a jinja2 friendly dict for all variants of the reduce API"""
    return ReduceConfig(
        scope=collectives_execution_scopes,
        all_dtypes=integral_dtypes + float_dtypes,
        integral_dtypes=integral_dtypes,
        reduction_op_literal=reduction_ops,
        bitwise_op_literal=bitwise_reduction_ops,
    ).to_dict()


def generate_coll(config) -> str:
    """Load the jinja2 template and render it with the config.
    
    Config Format:
    {
        "sync_config": {
            "scope": list[ExecutionScope],
        },
        "sync_all_config": {
            "scope": list[ExecutionScope],
        },
        "barrier_config": {
            "scope": list[ExecutionScope],
        },
        "barrier_all_config": {
            "scope": list[ExecutionScope],
        },
        "reduce_config": {
            "scope": list[ExecutionScope],
            "all_dtypes": list[Dtype],
            "integral_dtypes": list[Dtype],
            "reduction_op_literal": list[ReduceOp],
            "bitwise_op_literal": list[ReduceOp],
            "float_op_literal": list[ReduceOp],
        },
    }

    Returns
    -------
    str: The rendered Collectives bindings
    """
    this_dir = os.path.dirname(os.path.abspath(__file__))
    coll_template_path = os.path.join(this_dir, "templates", "core", "device", "numba", "collective.py.j2")
    with open(coll_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 = {
        "sync_config": generate_sync(),
        "sync_all_config": generate_sync_all(),
        "barrier_config": generate_barrier(),
        "barrier_all_config": generate_barrier_all(),
        "reduce_config": generate_reduce(),
        "reducescatter_config": generate_reducescatter(),
        "fcollect_config": generate_fcollect(),
        "broadcast_config": generate_broadcast(),
        "alltoall_config": generate_alltoall(),
    }
    target = os.path.join(args.output_dir, "collective.py")
    with open(target, "w") as f:
        print(f"Generating Collectives bindings to {target}")
        f.write(generate_coll(config))