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

import os
import re
from ast_canopy import parse_declarations_from_source
# from numbast import bind_cxx_struct, bind_cxx_function, MemoryShimWriter

from numba import types, cuda
from numba.core.datamodel.models import PrimitiveModel
import numpy as np
import yaml

with open(f"{os.path.dirname(os.path.abspath(__file__))}/config_nvshmem.yml", "r") as f:
    config = yaml.safe_load(f)

# Use AST Canopy to parse a header into declarations
source = config["Entry Point"]

source_allow_list = config["File List"]

additional_includes = config["Clang Include Paths"]

decls = parse_declarations_from_source(source,
                                       source_allow_list,
                                       additional_includes=additional_includes,
                                       compute_capability=config["GPU Arch"][0],
                                       bypass_parse_error=True)

# Authoritative mapping of NVSHMEM types to CuTe Types
# From d
cpp_to_cute_type = {
    "int": "cutlass.Int32",
    "float": "cutlass.Float32",
    "double": "cutlass.Float64",
    "bool": "cutlass.Boolean",
    "half": "cutlass.Float16",
    "bfloat16": "cutlass.BFloat16",
    "long": "cutlass.Int64",
    "long long": "cutlass.Int64",
    "void": None,
    "char": "cutlass.Int8",
    "unsigned char": "cutlass.Uint8",
    "short": "cutlass.Int16",
    "unsigned short": "cutlass.Uint16",
    "__nv_bfloat16": "cutlass.BFloat16",
    "__half": "cutlass.Float16",
    "signed char": "cutlass.Int8",
    "unsigned int": "cutlass.Uint32",
    "unsigned long": "cutlass.Uint64",
    "unsigned long long": "cutlass.Uint64",
    "size_t": "cutlass.Uint64",
    "int8_t": "cutlass.Int8",
    "uint8_t": "cutlass.Uint8",
    "int16_t": "cutlass.Int16",
    "uint16_t": "cutlass.Uint16",
    "int32_t": "cutlass.Int32",
    "uint32_t": "cutlass.Uint32",
    "int64_t": "cutlass.Int64",
    "uint64_t": "cutlass.Uint64",
    "float16_t": "cutlass.Float16",
    "float32_t": "cutlass.Float32",
    "float64_t": "cutlass.Float64",
    "char_t": "cutlass.Int8",
    "unsigned char_t": "cutlass.Uint8",
    "uintptr_t": "cutlass.Pointer",
    "nvshmemx_smem_amount_t": "cutlass.Int32",
    "enum nvshmemx_smem_amount_t": "cutlass.Int32",
}

# CuTe doesn't expose double2 (and related) types; skip them entirely.
_UNSUPPORTED_CUTE_TYPES = {
    "double2",
}

# Note: The following cutlass types are currently not supported by NVSHMEM:
# [SymInt, sym_int, sym_int32, sym_int64, Int4, Float8(all variants), TFloat32, ScaledBasis, IntTuple, Coord, Shape, Stride, Layout, Composedlayout,
# Tile, Tiler, XTuple]
# The following types are supported via type inference but will never be returned directly by the API:
# [Numeric, Integer]

_VOID_SIZE_TO_CUTE = {
    "8": "cutlass.Int8",
    "16": "cutlass.Int16",
    "32": "cutlass.Int32",
    "64": "cutlass.Int64",
}

_UNSUPPORTED_TYPES = {"double2"}


def _void_fallback(function_name: str) -> str:
    match = re.search(r"(\d+)$", function_name or "")
    if match:
        mapped = _VOID_SIZE_TO_CUTE.get(match.group(1))
        if mapped:
            return mapped
    return "cutlass.Int8"


def _strip_type_name(type_name: str) -> str:
    name = type_name.replace("const", "").replace("volatile", "").replace("&", "").strip()
    return name.split("*")[0].strip()


def _uses_unsupported_type(function) -> bool:
    type_names = [function.return_type.unqualified_non_ref_type_name]
    type_names.extend([param.type_.unqualified_non_ref_type_name for param in function.params])
    return any(_strip_type_name(name) in _UNSUPPORTED_TYPES for name in type_names)


def wrap_type_with_pointer(type_name, function_name=None, is_return=False):
    stripped_type = type_name.split("*")[0].strip()
    if stripped_type in _UNSUPPORTED_CUTE_TYPES:
        return None
    raw_type = cpp_to_cute_type.get(stripped_type, stripped_type)
    if raw_type is None:
        if stripped_type == "void" and is_return:
            return None
        if stripped_type == "void":
            raw_type = _void_fallback(function_name)
        else:
            raise RuntimeError(f"Unsupported NVSHMEM type in binding generation: '{stripped_type}'")
    if type_name.endswith("*"):
        return f"_CutePtrType({raw_type})"
    return raw_type


def generate_cute_decl(function, end="\n\n"):
    if _uses_unsupported_type(function):
        return None
    args_str = ", ".join([
        f"{wrap_type_with_pointer(arg.type_.unqualified_non_ref_type_name, function.name)}" for arg in function.params
    ])
    user_function_name = function.name.replace("nvshmem_", "").replace("nvshmemx_", "")
    return_type = wrap_type_with_pointer(
        function.return_type.unqualified_non_ref_type_name,
        function.name,
        is_return=True,
    )
    return_type_str = ""
    if return_type is not None:
        return_type_str = f", return_type={return_type}"
    return f"""{user_function_name} = cute.ffi(name=\"{function.name}\", params_types=[{args_str}]{return_type_str}){end}"""


prefix = """
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import cutlass
from cutlass import cute
from cutlass.cute.typing import AddressSpace

try:
    import cutlass._mlir.dialects.cute as _cute_ir
except Exception:  # pragma: no cover - depends on CuTe runtime
    _cute_ir = None

class _CutePtrType:
    def __init__(self, dtype, address_space=AddressSpace.gmem, alignment=None):
        self._dtype = dtype
        self._address_space = address_space
        self._alignment = alignment

    def __get_mlir_types__(self):
        if _cute_ir is None:
            raise RuntimeError("CuTe MLIR context is not available for pointer type generation")
        if self._alignment is None:
            align = getattr(self._dtype, 'width', 8) // 8 or 1
        else:
            align = self._alignment
        return [_cute_ir.PtrType.get(self._dtype.mlir_type, self._address_space, align)]
"""
all_function_decls = [prefix]
for function in decls.functions:
    function_decl = generate_cute_decl(function, end="\n")
    if function_decl is None:
        continue
    all_function_decls.append(function_decl)

print("Generated CuTe bindings for the following functions:")
for function in decls.functions:
    print(f"  {function.name}")

with open(f"{os.path.dirname(__file__)}/_cuteast.py", "w") as f:
    f.write("\n".join(all_function_decls))
