# Give functions object a name so we can metaprogram in the typed collectives
# Start with just hostlib init and hostlib_finalize
functions = {
    # Init/Fini
    'nvshmemx_hostlib_init_attr' : {},
    'nvshmemx_get_uniqueid' : {},
    'nvshmemx_set_attr_uniqueid_args' : {},
    'nvshmemx_set_attr_mpi_comm_args' : {},
    'nvshmemx_hostlib_finalize' : {
        "return": "TRANSPARENT",
    },
    'nvshmemx_cumodule_init': {
                "except?": -1,
                "return":"TRANSPARENT",
            },
    'nvshmemx_cumodule_finalize': {
        "except?": -1,
        "return":"TRANSPARENT",
    },
    'nvshmemx_culibrary_init': {
                "except?": -1,
                "return":"TRANSPARENT",
            },
    'nvshmemx_culibrary_finalize': {
        "except?": -1,
        "return":"TRANSPARENT",
    },
    # Memory Management
    'nvshmem_malloc' : {
        "return": "TRANSPARENT",
        "except?": 0,
    },
    'nvshmem_free' : {
        "return": "TRANSPARENT",
    },
    'nvshmem_calloc' : {
        'arg_names' : ("count", "size"),
        "return": "TRANSPARENT",
        "except?": 0,
    },
    'nvshmem_align' : {
        'arg_names' : ("count", "size"),
        "return": "TRANSPARENT",
        "except?": 0
    },
    'nvshmem_ptr' : {
        "return": "TRANSPARENT",
        "except?": 0,
        'arg_names' : ("dest", "pe"),
    },
    'nvshmemx_mc_ptr' : {
        "return": "TRANSPARENT",
        "except?": 0,
        'arg_names' : ("team", "ptr"),
    },
    'nvshmemx_buffer_register_symmetric' : {
        "return": "TRANSPARENT",
        "except?": 0,
    },
    'nvshmemx_buffer_unregister_symmetric' : {
        "return": "TRANSPARENT",
        "except?": 0,
    },
    # Team Management
    'nvshmem_my_pe' : {
        "return": "TRANSPARENT",
        "except?": -1,
    },
    'nvshmem_n_pes' : {
        "return": "TRANSPARENT",
        "except?": -1,
    },
    'nvshmem_team_my_pe' : {
        "return": "TRANSPARENT",
        "except?": -1,
    },
    'nvshmem_team_n_pes' : {
        "return": "TRANSPARENT",
        "except?": -1,
    },
    # RMA
    'nvshmemx_putmem_on_stream' : {
        "return": "TRANSPARENT",
    },
    'nvshmemx_getmem_on_stream' : {
        "return": "TRANSPARENT",
    },
    'nvshmemx_putmem_signal_on_stream' : {
        "return": "TRANSPARENT",
    },
    'nvshmemx_signal_op_on_stream' : {
        "return": "TRANSPARENT",
    },
    # Synchronization
    'nvshmem_barrier' : {},
    'nvshmem_barrier_all' : {"return": "TRANSPARENT",},
    'nvshmemx_barrier_on_stream' : {"except?": -1},
    'nvshmemx_team_sync_on_stream': {"return": "TRANSPARENT",  'except?': 0},
    'nvshmemx_barrier_all_on_stream' : {
                'arg_names' : ["stream"],
                "return": "TRANSPARENT",
            },
    'nvshmemx_sync_all_on_stream' : {
                'arg_names' : ["stream"],
                "return": "TRANSPARENT",
            },
    'nvshmemx_signal_wait_until_on_stream' : {
        "return": "TRANSPARENT",
    },
    'nvshmemx_quiet_on_stream': {
        "return": "TRANSPARENT",
    },
    # Misc.
    'nvshmem_info_get_version' : {"return": "TRANSPARENT",},
    'nvshmemx_vendor_get_version_info': {"return": "TRANSPARENT"},
    'nvshmemx_init_status' : {"return": "TRANSPARENT", "except?": 0},
    # [Arbitrary] team creation
    'nvshmem_team_get_config': {"return": "TRANSPARENT",},
    'nvshmem_team_translate_pe': {},
    'nvshmem_team_split_strided': {"except?": -1},
    'nvshmem_team_split_2d': {"except?": -1},
    'nvshmem_team_destroy': {"return": "TRANSPARENT",},
    'nvshmemx_team_init': {"except?": -1},
    'nvshmemx_team_get_uniqueid': {}
}

# Metaprogramming for all the various collectives
types = [
    # Floats
    'float', 'half', 'double', 'bfloat16',
    # UInts
    'uint8', 'uint16', 'uint32','uint64',
    # Ints
    'int8', 'short', 'int16', 'int32', 'int', 'int64', 'long', 'longlong',
    # Misc
    'size', 'char', 'schar',
]
operators = ["sum", "min", "max"]
collectives = {
    "fcollect": {"prefix": "nvshmemx", "need_op": False},
    "alltoall": {"prefix": "nvshmemx", "need_op": False},
    "reduce": {"prefix": "nvshmemx", "need_op": True},
    "reducescatter": {"prefix": "nvshmemx", "need_op": True},
    "broadcast": {"prefix": "nvshmemx", "need_op": False},
}
for type in types:
    for collective in collectives.keys():
        prefix = collectives[collective]["prefix"]
        need_op = collectives[collective]["need_op"]
        if not need_op:
            functions[f"{prefix}_{type}_{collective}_on_stream"] = {"except?": -1}
        else:
            for op in operators:
                functions[f"{prefix}_{type}_{op}_{collective}_on_stream"] = {"except?": -1,}  

config = {
    'nvshmem': {
        # lowpp module, the rest is auto-computed based on the layout rule
        'module': 'nvshmem.bindings.nvshmem',
        'data': {
            'versions': [
                # This is kind of a hack to make Cybind work
                # Call whatever is ToT "0.0.0" to avoid interfering with
                # other versions in Cybind's nvshmem dir
                ('3.3.0', ),
            ],
            # Only include host headers to avoid CPP parsing issues
            'headers': ['nvshmem_host.h'],
            'need_headers_at_build': False,
            # see the docstring of process_prefix_suffix() for the expected patterns
            'patterns': {
                'function': r'(^nvshmem)_(.*)|(^nvshmem)x_(.*)',
                'type': r'(^nvshmem)_(.*)_t$|(^nvshmem)_(?!team_config_v1)(.*)_v\d$|(^nvshmem)x_(.*)_t$|(^nvshmem)x_(?!init_args_v1$|init_attr_v1$)(.*)_v\d$|()(flags)|()(Cybind_enum\d+)|(^nvshmemx)_(status)',
            },
            # whether parsing the library headers requires also including CUDA headers
            'need_cuda': True,
            # whether checking the driver version is needed before loading the DSO/DLL
            'need_driver': False,
            # whether a Windows module is needed
            'support_win': False,
            # source of docstrings
            'docstrings': None,
        },
        'functions': functions,
        # use this to control the type apprearance at the lowpp level
        # this is used as WAR for current codegen limitations
        'types': {
            'nvshmemx_uniqueid_v1' : "AUTO_LOWPP",
            'nvshmem_team_config_t': "AUTO_LOWPP",
            'nvshmemx_team_uniqueid_t': "AUTO_LOWPP"
        },
        # map the enum values to their expected dtypes
        # this is very library-specific and needs the library developer to fill in
        # (or by referring to the library manual)
        'attrs': {
        },
        # This is used to patch (processed) enums, in case the generated lowpp enumerator names
        # are not meeting the expectation.
        'enums': {
            "^NVSHMEM_": "",
            "^NVSHMEMX_": "",
            "^NVSHMEMI_": "",
            "^PROXY_": "_PROXY_",
        },
    },
}

import pprint
print(pprint.pformat((config)))
