"""Python tree-sitter parser and function collection."""

from __future__ import annotations

import tree_sitter_python as tspython
from tree_sitter import Language, Node

from vibelexity.metrics.fatness.python import param_count
from vibelexity.parsers.shared import (
    collect_functions_wrapper,
    func_param_count_wrapper,
    parse_with_language,
)

PY_LANGUAGE = Language(tspython.language())

parse = lambda code: parse_with_language(PY_LANGUAGE, code)


def _collect_functions_recursive(node: Node, acc: list[Node]) -> None:
    if node.type == "function_definition":
        acc.append(node)
    for child in node.children:
        _collect_functions_recursive(child, acc)


def func_name(node: Node) -> str:
    """Extract function name from a function_definition node."""
    for child in node.children:
        if child.type == "identifier":
            return child.text.decode()
    return "<anonymous>"


collect_functions = collect_functions_wrapper(_collect_functions_recursive)
func_param_count = func_param_count_wrapper(param_count)
