#!/usr/bin/env python3

import argparse
import ast
import os
import sys


def parse_module(mod):
    with open(mod, "r") as source:
        tree = ast.parse(source.read())
    functions = [
        node
        for node in tree.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
    ]
    func_list = []
    for fun in functions:
        doc = ast.get_docstring(fun)
        doc = doc.splitlines()[0] if doc else ""
        func_list.append((fun.name, ast.unparse(fun.args), doc))

    classes = [node for node in tree.body if isinstance(node, ast.ClassDef)]
    class_list = []
    for cls in classes:
        doc = ast.get_docstring(cls)
        doc = doc.splitlines()[0] if doc else ""
        class_list.append(
            (cls.name, ", ".join([ast.unparse(base) for base in cls.bases]), doc)
        )

    assignments = [node for node in tree.body if isinstance(node, ast.Assign)]
    assign_list = []
    for assign in assignments:
        if isinstance(assign.value, ast.Call) and ast.unparse(
            assign.value.func
        ).endswith("namespaced_function"):
            pass  # We will get to that later
        else:
            tgt = assign.targets[0]
            tgtsrc = ast.unparse(tgt)
            if ((isinstance(tgt, ast.Name) and not tgtsrc.isupper())
                or (isinstance(tgt, ast.Subscript) and not tgt.value.id.isupper())):
                assign_list.append(
                    (tgtsrc, ast.unparse(assign.value))
                )

    for node in ast.walk(tree):
        if (
            isinstance(node, ast.Assign)
            and isinstance(node.value, ast.Call)
            and ast.unparse(node.value.func).endswith("namespaced_function")
        ):
            func_list.append((node.targets[0].id, "...", "NAMESPACED FUNCTION"))

    class_list.sort()
    func_list.sort()
    assign_list.sort()
    return class_list, func_list, assign_list


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "module",
        nargs="+",
        type=str,
        help="Path to Salt module",
    )
    options = parser.parse_args()

    first = True
    for mod in options.module:
        if not os.path.isfile(mod):
            parser.exit("Module should be a file: {}".format(mod))

        if first:
            first = False
        else:
            print("\n\n")
        print("### " + os.path.basename(mod))

        class_list, func_list, assign_list = parse_module(mod)
        for assign in assign_list:
            print(f"{assign[0]} = {assign[1]}")

        for cls in class_list:
            print(f"class {cls[0]}({cls[1]}): {cls[2]}")

        for fn in func_list:
            print(f"def {fn[0]}({fn[1]}): {fn[2]}")