"""Function and class definition extraction for tree-sitter.

Extracts function, class, and method definitions from JavaScript
and TypeScript parse trees.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from unfamiliar.analysis.adapters._treesitter_utils import get_node_text, traverse
from unfamiliar.analysis.models import FunctionDef

if TYPE_CHECKING:
    from tree_sitter import Node


def extract_definitions(root: Node) -> list[FunctionDef]:
    """Extract function, class, and method definitions from a tree-sitter tree.

    Captures:
    - function_declaration
    - generator_function_declaration
    - class_declaration
    - abstract_class_declaration (TS)
    - method_definition (inside class_body)
    - arrow_function inside variable_declarator

    Args:
        root: The root node of the parsed JS/TS file.

    Returns:
        List of FunctionDef objects for each definition.
    """
    definitions: list[FunctionDef] = []

    for node in traverse(root):
        if node.type in ("class_declaration", "abstract_class_declaration"):
            name_node = node.child_by_field_name("name")
            if name_node:
                class_name = get_node_text(name_node)

                definitions.append(
                    FunctionDef(
                        name=class_name,
                        line_start=node.start_point.row + 1,
                        line_end=node.end_point.row + 1,
                        is_async=False,
                        kind="class",
                    )
                )

                # Extract methods from class body
                body = node.child_by_field_name("body")
                if body:
                    for child in body.children:
                        if child.type == "method_definition":
                            method_def = _extract_method(child, class_name)
                            if method_def:
                                definitions.append(method_def)

        elif node.type in ("function_declaration", "generator_function_declaration"):
            name_node = node.child_by_field_name("name")
            if name_node:
                func_name = get_node_text(name_node)
                is_async = _has_async_modifier(node)

                definitions.append(
                    FunctionDef(
                        name=func_name,
                        line_start=node.start_point.row + 1,
                        line_end=node.end_point.row + 1,
                        is_async=is_async,
                        kind="function",
                    )
                )

        elif node.type == "variable_declarator":
            # Arrow function: const foo = () => {}
            name_node = node.child_by_field_name("name")
            value_node = node.child_by_field_name("value")

            if name_node and value_node and value_node.type == "arrow_function":
                var_name = get_node_text(name_node)
                is_async = _has_async_modifier(value_node)

                definitions.append(
                    FunctionDef(
                        name=var_name,
                        line_start=node.start_point.row + 1,
                        line_end=node.end_point.row + 1,
                        is_async=is_async,
                        kind="function",
                    )
                )

    return definitions


def _extract_method(node: Node, class_name: str) -> FunctionDef | None:
    """Extract a method definition from a method_definition node."""
    name_node = node.child_by_field_name("name")
    if not name_node:
        return None

    method_name = get_node_text(name_node)
    is_async = _has_async_modifier(node)

    return FunctionDef(
        name=f"{class_name}.{method_name}",
        line_start=node.start_point.row + 1,
        line_end=node.end_point.row + 1,
        is_async=is_async,
        kind="method",
    )


def _has_async_modifier(node: Node) -> bool:
    """Check if a function/method has the async modifier."""
    return any(child.type == "async" for child in node.children)
