# 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, collectives_execution_scopes, integral_dtypes as numbast_integral_dtypes,
    float_dtypes as numbast_float_dtypes, reduction_ops, bitwise_reduction_ops, ExecutionScope, BaseConfig, ReduceOp,
)

_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,
        }


def _convert_dtypes(dtypes):
    return [
        Dtype(
            numba_dtype_name=dtype.numba_dtype_name,
            nvshmem_dtype_name=dtype.nvshmem_dtype_name,
            cute_type_name=_CUTE_TYPE_NAMES[dtype.numba_dtype_name],
        ) for dtype in dtypes
    ]


integral_dtypes = _convert_dtypes(numbast_integral_dtypes)
float_dtypes = _convert_dtypes(numbast_float_dtypes)
all_dtypes = integral_dtypes + float_dtypes


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


@dataclass
class SyncConfig(CollectivesConfig):
    scope: list[ExecutionScope]


@dataclass
class SyncAllConfig(CollectivesConfig):
    scope: list[ExecutionScope]


@dataclass
class BarrierConfig(CollectivesConfig):
    scope: list[ExecutionScope]


@dataclass
class BarrierAllConfig(CollectivesConfig):
    scope: list[ExecutionScope]


@dataclass
class ReduceConfig(CollectivesConfig):
    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):
    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):
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]


@dataclass
class BroadcastConfig(CollectivesConfig):
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]


@dataclass
class AlltoallConfig(CollectivesConfig):
    scope: list[ExecutionScope]
    all_dtypes: list[Dtype]


def generate_sync() -> dict:
    return SyncConfig(scope=collectives_execution_scopes).to_dict()


def generate_sync_all() -> dict:
    return SyncAllConfig(scope=collectives_execution_scopes).to_dict()


def generate_barrier() -> dict:
    return BarrierConfig(scope=collectives_execution_scopes).to_dict()


def generate_barrier_all() -> dict:
    return BarrierAllConfig(scope=collectives_execution_scopes).to_dict()


def generate_reduce() -> dict:
    return ReduceConfig(
        scope=collectives_execution_scopes,
        all_dtypes=all_dtypes,
        integral_dtypes=integral_dtypes,
        reduction_op_literal=reduction_ops,
        bitwise_op_literal=bitwise_reduction_ops,
    ).to_dict()


def generate_reducescatter() -> dict:
    return ReducescatterConfig(
        scope=collectives_execution_scopes,
        all_dtypes=all_dtypes,
        integral_dtypes=integral_dtypes,
        reduction_op_literal=reduction_ops,
        bitwise_op_literal=bitwise_reduction_ops,
    ).to_dict()


def generate_fcollect() -> dict:
    return FCollectConfig(scope=collectives_execution_scopes, all_dtypes=all_dtypes).to_dict()


def generate_broadcast() -> dict:
    return BroadcastConfig(scope=collectives_execution_scopes, all_dtypes=all_dtypes).to_dict()


def generate_alltoall() -> dict:
    return AlltoallConfig(scope=collectives_execution_scopes, all_dtypes=all_dtypes).to_dict()


def generate_collective(config) -> str:
    template_path = os.path.join(this_dir, "templates", "core", "device", "cute", "collective.py.j2")
    with open(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 CuTe collective bindings to {target}")
        f.write(generate_collective(config))
