# %%
from typing import Callable
import ast
from docstring_parser import parse
from pathlib import Path

# %%

def check_non_existing_params(dp_names: list[str], sp_names: list[str], *, has_args, has_kwargs, ctx: str):
    for dp in dp_names:
        if dp not in sp_names and not (has_args or has_kwargs):
            print(f"E100 {ctx}: Argument '{dp}' found in docstring does not exist in function signature.")

def check_params_not_in_docstring(dp_names: list[str], sp_names: list[str], *, has_args, has_kwargs, ctx: str):
    for sp in sp_names:
        if sp not in dp_names and not (has_args or has_kwargs):
            print(f"E101 {ctx}: Argument '{sp}' found in signature but not in docstring.")


def compare_args(docstring_params, src_params, *, has_args, has_kwargs, ctx: str) -> bool:
    if src_params[0].arg in ("self", "cls"):
        src_params = src_params[1:]

    # Check ensemble of arguments
    dp_names = []
    for dp in docstring_params:
        dp_names.extend(_.strip() for _ in dp.arg_name.replace("*", "").split(","))

    sp_names = [sp.arg for sp in src_params]
    check_non_existing_params(dp_names, sp_names, has_args=has_args, has_kwargs=has_kwargs, ctx=ctx)
    # check_params_not_in_docstring(dp_names, sp_names, has_args=has_args, has_kwargs=has_kwargs, ctx=ctx)


def walk_ast_helper(path: Path) -> str:
    src = path.read_text()
    lines = src.splitlines()
    newLines = lines.copy()

    # Extract all functions and classes
    tree = ast.parse(src)
    nodes = [
        node
        for node in ast.walk(tree)
        if isinstance(node, (ast.ClassDef, ast.FunctionDef))
    ]

    # Iterate over docstrings in reversed order so that lines
    # can be modified
    for node in sorted(
        nodes, key=lambda node: node.body[0].lineno if node.body else 0
    ):
        docstring = ast.get_docstring(node)
        if not docstring:
            continue
        doc = parse(docstring)

        # For classes, we need to find the __init__ function down there
        args = []
        if isinstance(node, ast.ClassDef):
            for child_node in node.body:
                if isinstance(child_node, ast.FunctionDef) and child_node.name == "__init__":
                    args = child_node.args.posonlyargs + child_node.args.args + child_node.args.kwonlyargs
                    has_kwargs = True if child_node.args.kwarg else False
                    has_args = True if child_node.args.vararg else False
                    break
        else:
            args = node.args.posonlyargs + node.args.args + node.args.kwonlyargs
            has_kwargs = True if node.args.kwarg else False
            has_args = True if node.args.vararg else False

        if not doc.params or not args:
            continue
        ctx = f"{path}:{node.lineno}"
        compare_args(doc.params, args, has_kwargs=has_kwargs, has_args=has_args, ctx=ctx)


# %%
files = list((Path(__file__).parent / ".." ).glob("**/*.py"))
print(f"Found {len(files)} files")

for f in files:
    walk_ast_helper(f)

# %%