"""
Usage: module_name output_hpp input_dat[, input_dat...]

Generates a header file that contains initialization functions for pybind11 bindings

You must include the header "autogen_module_init.hpp", and call initWrapper() from
your pybind11 module declaration.
"""

import pathlib
import pickle
import sys
import typing as T

import toposort

from ..autowrap.buffer import RenderBuffer
from ..autowrap.context import HeaderContext


def _write_wrapper_hpp(
    module_name: str, output_hpp: pathlib.Path, *input_dat: pathlib.Path
):
    # Need to ensure that wrapper initialization is called in base order
    # so we have to toposort it here based on the class hierarchy determined
    # while parsing the header
    types2name = {}
    types2deps = {}
    ordering = []

    for datfile in input_dat:
        with open(datfile, "rb") as fp:
            hctx = pickle.load(fp)

        assert isinstance(hctx, HeaderContext)

        name = hctx.hname
        dep = hctx.class_hierarchy

        # make sure objects without classes are also included!
        if not dep:
            ordering.append(name)

        for clsname, bases in dep.items():
            if clsname in types2name:
                raise ValueError(f"{name} ({input_dat}): duplicate class {clsname}")
            types2name[clsname] = name
            types2deps[clsname] = bases[:]

    to_sort: T.Dict[str, T.Set[str]] = {}
    for clsname, bases in types2deps.items():
        clsname = types2name[clsname]
        deps = to_sort.setdefault(clsname, set())
        for base in bases:
            base = types2name.get(base)
            if base and base != clsname:
                deps.add(base)

    ordering.extend(toposort.toposort_flatten(to_sort, sort=True))

    r = RenderBuffer()
    r.writeln("// This file is autogenerated, DO NOT EDIT")
    r.writeln("")
    r.write_trim(
        f"""
        #pragma once
        #include <semiwrap.h>

        // Use this to define your module instead of PYBIND11_MODULE
        #define SEMIWRAP_PYBIND11_MODULE(variable) PYBIND11_MODULE({module_name}, variable)

        // TODO: namespace semiwrap::autogen {{

        // forward declarations
        """
    )
    for name in ordering:
        r.write_trim(
            f"""
                void begin_init_{name}(py::module &m);
                void finish_init_{name}();
            """
        )
        r.writeln()

    r.writeln("static void initWrapper(py::module &m) {")
    with r.indent():
        for name in ordering:
            r.writeln(f"begin_init_{name}(m);")
        r.writeln()
        for name in ordering:
            r.writeln(f"finish_init_{name}();")

    r.writeln("}")

    output_hpp.write_text(r.getvalue(), encoding="utf-8")


def main():
    try:
        module_name = sys.argv[1]
        output_hpp = sys.argv[2]
        inputs = sys.argv[3:]
    except Exception as e:
        print(__doc__, file=sys.stderr)
        sys.exit(1)

    _write_wrapper_hpp(
        module_name, pathlib.Path(output_hpp), *[pathlib.Path(i) for i in inputs]
    )


if __name__ == "__main__":
    main()
