"""
This pre-processor parses a single file containing a list of
`MP_REGISTER_MODULE(MP_QSTR_module_name, obj_module)` or
`MP_REGISTER_EXTENSIBLE_MODULE(MP_QSTR_module_name, obj_module)`
(i.e. the output of `py/makeqstrdefs.py cat module`).

The output is a header (typically moduledefs.h) which is included by
py/objmodule.c that contains entries to be included in the definition of
 - mp_rom_map_elem_t mp_builtin_module_table[]
 - mp_rom_map_elem_t mp_builtin_extensible_module_table[]

Extensible modules are modules that can be overridden from the filesystem, see
py/builtinimnport.c:process_import_at_level. Regular modules will always use
the built-in version.
"""

import sys
import re
import io
import argparse


register_pattern = re.compile(
    r"\s*(MP_REGISTER_MODULE|MP_REGISTER_EXTENSIBLE_MODULE)\(MP_QSTR_(.*?),\s*(.*?)\);",
    flags=re.DOTALL,
)

delegation_pattern = re.compile(
    r"\s*(?:MP_REGISTER_MODULE_DELEGATION)\((.*?),\s*(.*?)\);",
    flags=re.DOTALL,
)


def find_module_registrations(filename):
    """Find any MP_REGISTER_MODULE definitions in the provided file.

    :param str filename: path to file to check
    :return: List[(module_name, obj_module)]
    """
    global pattern

    with io.open(filename, encoding="utf-8") as c_file_obj:
        c = c_file_obj.read()
        return set(re.findall(register_pattern, c)), set(re.findall(delegation_pattern, c))


def generate_module_table_header(modules):
    """Generate header with module table entries for builtin modules.

    :param List[(module_name, obj_module)] modules: module defs
    :return: None
    """

    # Print header file for all external modules.
    mod_defs = set()
    extensible_mod_defs = set()
    for macro_name, module_name, obj_module in modules:
        mod_def = "MODULE_DEF_{}".format(module_name.upper())
        if macro_name == "MP_REGISTER_MODULE":
            mod_defs.add(mod_def)
        elif macro_name == "MP_REGISTER_EXTENSIBLE_MODULE":
            extensible_mod_defs.add(mod_def)
        if "," in obj_module:
            print(
                "ERROR: Call to {}({}, {}) should be {}({}, {})\n".format(
                    macro_name,
                    module_name,
                    obj_module,
                    macro_name,
                    module_name,
                    obj_module.split(",")[0],
                ),
                file=sys.stderr,
            )
            sys.exit(1)
        print(
            (
                "extern const struct _mp_obj_module_t {obj_module};\n"
                "#undef {mod_def}\n"
                "#define {mod_def} {{ MP_ROM_QSTR(MP_QSTR_{module_name}), MP_ROM_PTR(&{obj_module}) }},\n"
            ).format(
                module_name=module_name,
                obj_module=obj_module,
                mod_def=mod_def,
            )
        )

    # There should always be at least one module (__main__ in runtime.c)
    assert mod_defs

    print("\n#define MICROPY_REGISTERED_MODULES \\")

    for mod_def in sorted(mod_defs):
        print("    {mod_def} \\".format(mod_def=mod_def))
    print("// MICROPY_REGISTERED_MODULES")

    # There are not necessarily any extensible modules (e.g., bare-arm or minimal x86)
    print("\n#define MICROPY_HAVE_REGISTERED_EXTENSIBLE_MODULES ", len(extensible_mod_defs))

    if extensible_mod_defs:
        print("\n#define MICROPY_REGISTERED_EXTENSIBLE_MODULES \\")

        for mod_def in sorted(extensible_mod_defs):
            print("    {mod_def} \\".format(mod_def=mod_def))

        print("// MICROPY_REGISTERED_EXTENSIBLE_MODULES")


def generate_module_delegations(delegations):
    if not delegations:
        return

    print()
    for obj_module, fun_name in delegations:
        print("extern void {}(mp_obj_t self_in, qstr attr, mp_obj_t *dest);".format(fun_name))
    print("#define MICROPY_MODULE_DELEGATIONS \\")
    for obj_module, fun_name in delegations:
        print(
            "    {{ MP_ROM_PTR(&{obj_module}), {fun_name} }}, \\".format(
                obj_module=obj_module, fun_name=fun_name
            )
        )
    print("// MICROPY_MODULE_DELEGATIONS")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("file", nargs=1, help="file with MP_REGISTER_MODULE definitions")
    args = parser.parse_args()

    print("// Automatically generated by makemoduledefs.py.\n")

    modules, delegations = find_module_registrations(args.file[0])
    generate_module_table_header(sorted(modules))
    generate_module_delegations(sorted(delegations))


if __name__ == "__main__":
    main()
