"""Data Transformation Density (DTD) for TypeScript."""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from tree_sitter import Node

_TS_DTD_CONTAINERS = {"object", "array"}

_TS_DTD_PAYLOADS = {"pair", "shorthand_property_identifier"}

_TS_DTD_MAP_NAMES = {"map", "filter", "reduce", "flatMap", "forEach"}

_TS_DTD_PUNCTUATION = {"[", "]", "{", "}", ",", "comment"}


def compute_dtd_ts(func_node: Node) -> int:
    """Compute Data Transformation Density for a TypeScript function node."""
    for child in func_node.children:
        if child.type == "statement_block":
            return _dtd_walk_ts(child, depth=0, top_func=func_node)
    if func_node.type == "arrow_function":
        children = [
            c
            for c in func_node.children
            if c.type not in ("=>", "formal_parameters", "identifier")
        ]
        for child in children:
            return _dtd_walk_ts(child, depth=0, top_func=func_node)
    return 0


def _is_template_container(node: Node) -> bool:
    """Check if a template_string has template_substitution children (making it a container)."""
    if node.type != "template_string":
        return False
    return any(c.type == "template_substitution" for c in node.children)


def _is_map_call_ts(node: Node) -> bool:
    """Check if a call_expression is a map/filter/reduce/flatMap/forEach call."""
    if node.type != "call_expression":
        return False
    if not node.children:
        return False
    first = node.children[0]
    if first.type == "member_expression":
        prop = first.child_by_field_name("property")
        if prop and prop.text.decode() in _TS_DTD_MAP_NAMES:
            return True
    return False


_TS_DTD_NESTED_FUNCS = {
    "function_declaration",
    "function_expression",
    "method_definition",
    "generator_function_declaration",
}


def _dtd_is_nested_func(child: Node, top_func: Node) -> bool:
    """Check if a child is a nested function that should be skipped or handled specially."""
    if child.type in _TS_DTD_NESTED_FUNCS and child is not top_func:
        return True
    if child.type == "arrow_function" and child is not top_func:
        return True
    return False


def _dtd_increases_depth(child) -> bool:
    """Check if a child node should recurse with depth + 1."""
    if child.type in _TS_DTD_CONTAINERS:
        return True
    if _is_template_container(child):
        return True
    if _is_map_call_ts(child):
        return True
    return False


def _dtd_is_payload(child: Node, node: Node) -> bool:
    """Check if a child is a payload node that scores at current depth."""
    if child.type in _TS_DTD_PAYLOADS:
        return True
    if child.type == "template_substitution" and node.type == "template_string":
        return True
    return False


def _dtd_walk_ts(node: Node, depth: int, top_func: Node) -> int:
    """Walk AST and accumulate DTD score."""
    total = 0
    is_array = node.type == "array"
    for child in node.children:
        if _dtd_is_nested_func(child, top_func):
            if child.type == "arrow_function" and not (
                child.parent and child.parent.type == "variable_declarator"
            ):
                total += _dtd_walk_ts(child, depth, top_func)
            continue

        is_element = is_array and child.type not in _TS_DTD_PUNCTUATION
        if is_element:
            total += depth

        if _dtd_increases_depth(child):
            total += _dtd_walk_ts(child, depth + 1, top_func)
        elif _dtd_is_payload(child, node):
            if not is_element:
                total += depth
            total += _dtd_walk_ts(child, depth, top_func)
        else:
            total += _dtd_walk_ts(child, depth, top_func)

    return total
