from __future__ import annotations

import ast
from pathlib import Path


ANZATZ_DIR = Path(__file__).resolve().parents[1] / "src" / "topovec" / "ansatz"
MODULES = [
    "__init__.py",
    "_utils.py",
    "base.py",
    "field.py",
    "frame.py",
    "mask.py",
    "section.py",
]


def _load_module_tree(name: str) -> ast.Module:
    return ast.parse((ANZATZ_DIR / name).read_text())


def _public_top_level_defs(tree: ast.Module):
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and not node.name.startswith("_"):
            yield node


def _public_methods(class_node: ast.ClassDef):
    for node in class_node.body:
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            continue
        if node.name == "__init__":
            continue
        if node.name == "__call__" or not node.name.startswith("_"):
            yield node


def test_ansatz_modules_have_docstrings():
    for name in MODULES:
        tree = _load_module_tree(name)
        assert ast.get_docstring(tree), f"{name} is missing a module docstring."


def test_public_ansatz_module_definitions_have_docstrings():
    for name in MODULES:
        tree = _load_module_tree(name)
        for node in _public_top_level_defs(tree):
            assert ast.get_docstring(node), f"{name}:{node.name} is missing a public docstring."


def test_public_ansatz_methods_have_docstrings():
    for name in MODULES:
        tree = _load_module_tree(name)
        for node in _public_top_level_defs(tree):
            if not isinstance(node, ast.ClassDef):
                continue
            for method in _public_methods(node):
                assert ast.get_docstring(method), (
                    f"{name}:{node.name}.{method.name} is missing a public method docstring."
                )
