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

from __future__ import annotations

import tree_sitter_typescript as tstypescript
from tree_sitter import Language, Node, Parser

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

TS_LANGUAGE = Language(tstypescript.language_typescript())
TSX_LANGUAGE = Language(tstypescript.language_tsx())


def parse(code: str, tsx: bool = False) -> Node:
    """Parse TypeScript source code and return the tree root node."""
    lang = TSX_LANGUAGE if tsx else TS_LANGUAGE
    parser = Parser(lang)
    tree = parser.parse(code.encode())
    return tree.root_node


def _collect_functions_recursive(node: Node, acc: list[Node]) -> None:
    func_types = {
        "function_declaration",
        "function_expression",
        "method_definition",
        "generator_function_declaration",
    }
    if node.type in func_types:
        acc.append(node)
    elif node.type == "arrow_function":
        if node.parent and node.parent.type == "variable_declarator":
            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 node."""
    if node.type == "arrow_function":
        if node.parent and node.parent.type == "variable_declarator":
            for child in node.parent.children:
                if child.type == "identifier":
                    return child.text.decode()
        return "<anonymous>"
    if node.type == "method_definition":
        for child in node.children:
            if child.type == "property_identifier":
                return child.text.decode()
        return "<anonymous>"
    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)
