"""CLI to generate a .pyi stub file from an extension module.

The stub definitions WILL include:
- All classes defined in the module.
- Immediate ancestors of each class in the inheritance list.
- All method descriptors in classes (with some exceptions).
- All data descriptors, i.e. getsets (`PyGetSetDef`) and members (`PyMemberDef`).
- Plain-old attributes (also a catch-all) ONLY if they do not simply evaluate to `None`.
  - These will have a deduced type annotation.
- Docstrings (where present in the extension module).
- Type annotations for methods ONLY if they are defined in the module.

The stub definitions WILL NOT include:
- Other module-level objects (functions etc.)
- Methods which are:
  - Recognised as built-in.
  - Explicitly ignored in `EXEMPT_MEMBERS`.
- Deduced parameter/return type annotations for functions.
- Deduced type annotations for data descriptors.

In addition, an `__iter__` stub will be added to classes which are old-style iterables, i.e.
classes which define `__getitem__` and `__len__` but not `__iter__`. This allows these classes
to be used as an `Iterable` by type checkers. This behaviour can be disabled in order to be more
honest with the source material.
"""

import inspect
import re
import argparse
import importlib
import io
import pathlib
import types


class FileWriter:
    INDENT_SPACES = 4

    def __init__(self, out: io.TextIOWrapper):
        self.out = out
        self._level = 0

    def indent(self) -> None:
        self._level += 1

    def dedent(self) -> None:
        if self._level == 0:
            raise RuntimeError("Cannot dedent below level 0.")
        self._level -= 1

    def write(self, s: str) -> None:
        # Remove any backslashes which do not form a valid escape sequence
        s_sanitised = s.replace("\\", "")

        # Count the number of line breaks in the string
        n_line_breaks = s_sanitised.count("\n")

        # If the string is ONLY line breaks, just print it now
        if len(s_sanitised) == n_line_breaks:
            self.out.write(s_sanitised)
        # Else, add correct indentation
        else:
            indentation = self._indent_str() * self._level

            # First add indentation to beginning of string
            s_indented = indentation + s_sanitised

            # Add indentation after all line breaks
            s_indented = s_indented.replace(
                "\n",
                "\n" + indentation,
            )

            # Remove unnecessary indentation from empty intermediate lines
            processing = True
            while processing:
                s_indented, count = re.subn(
                    r"\n +\n",
                    "\n\n",
                    s_indented
                )
                processing = count != 0

            # Remove unnecessary indentation from end of string, if present
            if len(s_indented) >= len(indentation) and s_indented[-len(indentation):] == indentation:
                s_indented = s_indented[:-len(indentation)]

            self.out.write(s_indented)

    @classmethod
    def _indent_str(cls) -> str:
        return " " * cls.INDENT_SPACES


EXEMPT_MEMBERS: list[str] = [
    "__doc__",
    "__module__",
    "__class__",
]


def make_stub_file(
        package: types.ModuleType,
        module: types.ModuleType,
        output_file: pathlib.Path,
        debug: bool = False,
        mock_iter: bool = True
    ) -> None:
    with open(output_file, "w") as file:
        f = FileWriter(file)

        if module.__doc__ is not None:
            f.write("\"\"\"")
            f.write(f"{module.__doc__}\n")
            f.write(f"THIS FILE WAS GENERATED BY make_stubs FROM {package.__name__} VERSION {package.__version__}.\n")
            f.write("\"\"\"\n\n")
        else:
            f.write(f"\"\"\"THIS FILE WAS GENERATED BY make_stubs FROM {package.__name__} VERSION {package.__version__}.\"\"\"\n\n")

        f.write("from _typeshed import Incomplete\n")

        for cls_name, cls in inspect.getmembers(module):
            if inspect.isclass(cls):
                f.write("\n\n")
                bases = cls.__bases__

                if len(bases) == 1 and bases[0] == object:
                    f.write(f"class {cls_name}:\n")
                else:
                    f.write(f"class {cls_name}(")
                    for base in bases:
                        f.write(base.__name__)
                    f.write("):\n")

                f.indent()
                f.write(f"\"\"\"{cls.__doc__}\"\"\"\n")

                for member_name, member in inspect.getmembers(cls):
                    if member_name not in EXEMPT_MEMBERS:
                        # Skip built-ins
                        if inspect.isbuiltin(member):
                            continue

                        f.write("\n")

                        # Methods
                        if inspect.ismethoddescriptor(member):
                            f.write(f"def {member_name}{inspect.signature(member)}:\n")
                            f.indent()
                            if member.__doc__ is not None and member.__doc__ != "":
                                f.write(f"\"\"\"{member.__doc__}\"\"\"\n")
                            f.write("...\n")
                            f.dedent()

                        # Data descriptor
                        elif inspect.isdatadescriptor(member):
                            f.write(f"{member_name}: Incomplete  # DATA DESCRIPTOR\n")
                            if member.__doc__ is not None and member.__doc__ != "":
                                f.write(f"\"\"\"{member.__doc__}\"\"\"\n")

                        # Plain-old attribute
                        else:
                            if member is not None:
                                f.write(f"{member_name}: {type(member).__name__}\n")

                        if debug:
                            f.write(f"# ismethoddescriptor = {inspect.ismethoddescriptor(member)}\n")
                            f.write(f"# ismethod = {inspect.ismethod(member)}\n")
                            f.write(f"# isfunction = {inspect.isfunction(member)}\n")
                            f.write(f"# isbuiltin = {inspect.isbuiltin(member)}\n")
                            f.write(f"# isdatadescriptor = {inspect.isdatadescriptor(member)}\n")
                            f.write(f"# isgetsetdescriptor = {inspect.isgetsetdescriptor(member)}\n")
                            f.write(f"# ismemberdescriptor = {inspect.ismemberdescriptor(member)}\n")
                            f.write(f"# isclass = {inspect.isclass(member)}\n")

                if mock_iter:
                    names: list[str] = [tup[0] for tup in inspect.getmembers(cls)]
                    if "__getitem__" in names and "__iter__" not in names:
                        f.write("\n")
                        f.write("def __iter__(self):\n")
                        f.indent()
                        f.write("\"\"\"make_stubs: MOCKED TO REGISTER THIS CLASS AS AN OLD-STYLE ITERABLE.\"\"\"\n")
                        f.write("...\n")
                        f.dedent()

                f.dedent()


def main() -> None:
    parser = argparse.ArgumentParser(
        prog="make_stubs",
        description="Generate a .pyi stub file from an extension module"
    )

    parser.add_argument("module", help="full name of extension module")
    parser.add_argument("output", help="root directory for output files")
    parser.add_argument("-d", "--debug", action="store_true", help="enable verbose debug comments in output")
    parser.add_argument("--no-iter", action="store_true", help="disable mocking __iter__ on old-style iterables")

    args = parser.parse_args()

    module = importlib.import_module(args.module)

    root_dir = pathlib.Path(args.output)
    module_tree = str(args.module).split(".")

    package = importlib.import_module(module_tree[0])

    output_file = root_dir
    for name in module_tree[:-1]:
        output_file /= name

    output_file /= (module_tree[-1] + ".pyi")
    output_file.parent.mkdir(parents=True, exist_ok=True)

    make_stub_file(package, module, output_file, args.debug, not args.no_iter)


if __name__ == "__main__":
    main()
