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

import os
import argparse

import jinja2

this_dir = os.path.dirname(os.path.abspath(__file__))

_CUTE_TYPE_NAMES = {
    "int32": "Int32",
    "int64": "Int64",
    "uint32": "Uint32",
    "uint64": "Uint64",
    "float32": "Float32",
    "float64": "Float64",
}

_NVSHMEM_TYPE_NAMES = {
    "int32": "int32",
    "int64": "long",  # CuTe bindings use "long" not "int64" for atomic operations
    "uint32": "uint32",
    "uint64": "uint64",  # Note: uint64 exists in bindings, but ulong also exists
    "float32": "float",
    "float64": "double",
}


class Dtype:

    def __init__(self, dtype_name):
        self.dtype_name = dtype_name
        self.nvshmem_type_name = _NVSHMEM_TYPE_NAMES[dtype_name]
        self.cute_type_name = _CUTE_TYPE_NAMES[dtype_name]

    def to_dict(self):
        return {
            "dtype_name": self.dtype_name,
            "nvshmem_type_name": self.nvshmem_type_name,
            "cute_type_name": self.cute_type_name,
        }


def _make_dtypes(names):
    return [Dtype(name).to_dict() for name in names]


def generate_amo() -> str:
    amo_template_path = os.path.join(this_dir, "templates", "core", "device", "cute", "amo.py.j2")
    with open(amo_template_path, "r") as f:
        template = jinja2.Template(f.read())

    int_dtypes = ["int32", "int64", "uint32", "uint64"]
    float_dtypes = ["float32", "float64"]

    config = {
        "amo_fetch_dtypes": _make_dtypes(int_dtypes + float_dtypes),
        "amo_set_dtypes": _make_dtypes(int_dtypes + float_dtypes),
        "amo_swap_dtypes": _make_dtypes(int_dtypes + float_dtypes),
        "amo_compare_swap_dtypes": _make_dtypes(int_dtypes),
        "amo_fetch_inc_dtypes": _make_dtypes(int_dtypes),
        "amo_inc_dtypes": _make_dtypes(int_dtypes),
        "amo_fetch_add_dtypes": _make_dtypes(int_dtypes),
        "amo_add_dtypes": _make_dtypes(int_dtypes),
        "amo_bitwise_dtypes": _make_dtypes(int_dtypes),
    }

    return template.render(**config)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=str, default=".")
    args = parser.parse_args()

    target = os.path.join(args.output_dir, "amo.py")
    with open(target, "w") as f:
        print(f"Generating CuTe AMO bindings to {target}")
        f.write(generate_amo())
