#! /usr/bin/env python
"""Python script to make a hugo compatible code documentation."""

import ast
import argparse
import dataclasses
import inspect
import re
import pathlib

import tree_sitter
import tree_sitter_python


__authors__ = ["Lucas ROUAUD"]
__contact__ = ["lucas.rouaud@gmail.com"]
__version__ = "0.0.1"
__date__ = "2026-02-11"
__copyright__ = ["MIT License", "CC-BY-SA"]


def check_path(directory_path: str) -> pathlib.Path:
    directory = pathlib.Path(directory_path).expanduser().resolve()

    if not (directory.exists() and directory.is_dir()):
        raise NotADirectoryError(
            f"[Err##] '{directory_path}' does not exists or is not a "
            + "directory."
        )

    return directory


class ArgumentTyping(argparse.Namespace):
    input: pathlib.Path = pathlib.Path()
    output: pathlib.Path = pathlib.Path()
    skip_underscored: bool = False


def argument() -> ArgumentTyping:
    parser = argparse.ArgumentParser(
        description=(
            "Python script used to generated a code documentation compatible "
            + "with Hugo."
        ),
        formatter_class=argparse.RawTextHelpFormatter,
        add_help=False,
    )

    # == REQUIRED.
    _ = parser.add_argument(
        "-i",
        "--input",
        dest="input",
        required=True,
        type=check_path,
        metavar="[DIRECTORY]",
        help="Module directory to make code documentation from.",
    )

    _ = parser.add_argument(
        "-o",
        "--output",
        dest="output",
        required=True,
        type=check_path,
        metavar="[DIRECTORY]",
        help="Directory to output the code documentation.",
    )

    # == OPTIONAL.
    _ = parser.add_argument(
        "-h",
        "--help",
        action="help",
        help="Display this help message, then exit.",
    )

    _ = parser.add_argument(
        "-v",
        "--version",
        action="version",
        version=__version__,
        help="Display the program version, then exit.",
    )

    _ = parser.add_argument(
        "-s",
        "--skip_underscored",
        dest="skip_underscored",
        action="store_true",
        help=(
            "If given, the program will skip all files containing double "
            + "underscore, like '__init__.py' or '__main__.py"
        ),
    )

    return parser.parse_args(namespace=ArgumentTyping())


def get_path_list(
    directory: pathlib.Path, skip_underscored: bool = False
) -> list[pathlib.Path]:
    python_file: list[pathlib.Path] = []

    for element in directory.iterdir():
        if element.is_dir():
            python_file.extend(get_path_list(element, skip_underscored))
            continue

        if element.suffix != ".py":
            continue

        file_name: str = element.stem

        if (
            skip_underscored
            and file_name.startswith("__")
            and file_name.endswith("__")
        ):
            continue

        python_file.append(element)

    return python_file


@dataclasses.dataclass
class Function:
    name: str
    docstring: str
    content: str
    decorator: list[str] = dataclasses.field(default_factory=list)
    type: str = "function"


@dataclasses.dataclass
class Class:
    name: str
    docstring: str
    content: str
    decorator: list[str] = dataclasses.field(default_factory=list)
    method: list[Function] = dataclasses.field(default_factory=list)
    type: str = "class"


@dataclasses.dataclass
class Module:
    name: str
    docstring: str
    type: str = "module"


class ModuleParser(list[Function | Class | Module]):
    def __init__(
        self, source: pathlib.Path, language: tree_sitter.Language
    ) -> None:
        self.source: bytes = source.read_bytes()
        parser: tree_sitter.Parser = tree_sitter.Parser(language)
        root_node: tree_sitter.Node = parser.parse(self.source).root_node

        self.append(
            Module(name=source.stem, docstring=self.parse_docstring(root_node))
        )

        for child in root_node.children:
            definition: Class | Function | None = None

            match child.type:
                case "function_definition" | "async_function_definition":
                    definition = self.parse_function(child)
                case "decorated_definition":
                    definition = self.parse_decorator(child)
                case "class_definition":
                    definition = self.parse_class(child)
                case _:
                    pass

            if definition is not None:
                self.append(definition)

    def get_text(self, node: tree_sitter.Node | None) -> str:
        if node is None:
            return ""

        return self.source[node.start_byte : node.end_byte].decode()

    def parse_docstring(self, node: tree_sitter.Node | None) -> str:
        if node is None:
            return ""

        statement: list[tree_sitter.Node] = []

        for child_node in node.children:
            if child_node.is_named:
                statement.append(child_node)

        if len(statement) <= 0:
            return ""

        first_node: tree_sitter.Node = statement[0]

        if first_node.type != "expression_statement":
            return ""

        if len(first_node.children) <= 0:
            return ""

        string_node: tree_sitter.Node = first_node.children[0]

        if string_node.type != "string":
            return ""

        return self.get_text(string_node)

    def parse_decorator(
        self, node: tree_sitter.Node
    ) -> Function | Class | None:
        decorator: list[str] = []

        for child in node.children:
            if child.type != "decorator":
                continue

            decorator.append(self.get_text(child))

        definition: tree_sitter.Node = node.children[-1]

        match definition.type:
            case "function_definition":
                return self.parse_function(definition, decorator)
            case "class_definition":
                return self.parse_class(definition, decorator)
            case _:
                pass

        return None

    def parse_function(
        self, node: tree_sitter.Node, decorator: list[str] | None = None
    ) -> Function:
        name: str = self.get_text(node.child_by_field_name("name"))
        content: str = self.get_text(node)
        docstring: str = self.parse_docstring(node.child_by_field_name("body"))

        return Function(
            name=name,
            docstring=docstring,
            content=content,
            decorator=decorator or [],
        )

    def parse_class(
        self, node: tree_sitter.Node, decorator: list[str] | None = None
    ) -> Class:
        name: str = self.get_text(node.child_by_field_name("name"))
        content: str = self.get_text(node)
        docstring: str = self.parse_docstring(node.child_by_field_name("body"))

        # Method parsing.
        method: list[Function] = []
        block: tree_sitter.Node | None = node.child_by_field_name("body")

        if block is None:
            return Class(
                name=name,
                docstring=docstring,
                content=content,
                method=method,
                decorator=decorator or [],
            )

        for child in block.children:
            if not child.is_named:
                continue

            match child.type:
                case "function_definition":
                    method.append(self.parse_function(child))
                case "decorated_definition":
                    result: Class | Function | None = self.parse_decorator(
                        child
                    )

                    if isinstance(result, Function):
                        method.append(result)
                case _:
                    pass

        return Class(
            name=name,
            docstring=docstring,
            content=content,
            method=method,
            decorator=decorator or [],
        )


class MarkdownWritter:
    def __init__(self, input: pathlib.Path, output: pathlib.Path) -> None:
        self.input: pathlib.Path = input
        self.output: pathlib.Path = output

    def __call__(self, module: ModuleParser, file: pathlib.Path) -> None:
        content: str = ""

        for element in module:
            if isinstance(element, Module):
                content += self.title(element.name)
                content += self.docstring(element.docstring)
                content += "{{< katex />}}\n\n"
            elif isinstance(element, Function):
                content += self.title(element.name, 2)
                content += self.docstring(element.docstring, 2)
                content += self.code_block(element.content)
            elif isinstance(element, Class):
                content += self.title(element.name, 2)
                content += self.docstring(element.docstring, 2)
                content += self.code_block(element.content)

                if element.method == []:
                    continue

                for sub_element in element.method:
                    content += self.title(sub_element.name, 3)
                    content += self.docstring(sub_element.docstring, 3)
                    content += self.code_block(sub_element.content)

        output_name: str = str(file.relative_to(self.input).with_suffix(".md"))
        output_name = output_name.replace("/", ".")

        with open(
            self.output / output_name, "w", encoding="utf-8"
        ) as markdown:
            _ = markdown.write(content)

    def title(self, content: str, level: int = 1) -> str:
        if level < 1:
            raise ValueError(
                "[Err##] 'level' must be greater or equal then 1."
            )

        return f"{'#' * level} {content.replace('_', r'\_')}\n\n"

    def docstring(self, docstring: str, current_level: int = 1) -> str:
        """
        Special format `%%` at the enf of a line == delete the `\n` + spaces if
        there are.
        """
        if docstring == "":
            return ""

        if current_level < 1:
            raise ValueError(
                "[Err##] 'current_level' must be greater or equal then 1."
            )

        value: str = ast.literal_eval(docstring)
        cleaned_docstring: str = inspect.cleandoc(value) + "\n\n"

        header: str = "#" * (current_level + 1)
        cleaned_header: str = re.sub(
            r"([^\n]+)\n-{3,}", rf"{header} \1\n", cleaned_docstring
        )
        cleaned_space: str = re.sub(
            r"\n([^\s]{4}.+)\n    ", r"\n**\1**\n\n    ", cleaned_header
        )
        special_format: str = re.sub(r"%%\n *", " ", cleaned_space)

        return special_format.replace("    ", "")

    def code_block(self, code: str) -> str:
        return (
            '{{% details title="Code block details" open=false %}}\n'
            + f"```python\n{code}\n```\n"
            + "{{% /details %}}\n"
        )


def main() -> None:
    parsed_argument: ArgumentTyping = argument()

    file_list: list[pathlib.Path] = get_path_list(
        parsed_argument.input, parsed_argument.skip_underscored
    )

    module: list[ModuleParser] = []
    language = tree_sitter.Language(tree_sitter_python.language())
    to_markdown = MarkdownWritter(
        input=parsed_argument.input, output=parsed_argument.output
    )

    for file in file_list:
        parsed_module = ModuleParser(file, language)
        to_markdown(parsed_module, file)


if __name__ == "__main__":
    main()
