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

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from tree_sitter import Node

_PY_DTD_CONTAINERS = {
    "dictionary",
    "list",
    "set",
    "list_comprehension",
    "dictionary_comprehension",
    "set_comprehension",
    "generator_expression",
}

_PY_DTD_PAYLOADS = {"pair", "keyword_argument"}

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

_PY_DTD_MAP_NAMES = {"map", "filter", "reduce"}


def compute_dtd_py(func_node: Node) -> int:
    """Compute Data Transformation Density for a Python function_definition node."""
    for child in func_node.children:
        if child.type == "block":
            return _dtd_walk_py(child, depth=0, top_func=func_node)
    return 0


def _is_map_call_py(node: Node) -> bool:
    """Check if a call node is a map/filter/reduce call."""
    if node.type != "call":
        return False
    if not node.children:
        return False
    first = node.children[0]
    if first.type == "attribute":
        for child in first.children:
            if child.type == "identifier" and child.text.decode() in _PY_DTD_MAP_NAMES:
                return True
    if first.type == "identifier" and first.text.decode() in _PY_DTD_MAP_NAMES:
        return True
    return False


def _dtd_walk_py(node: Node, depth: int, top_func: Node) -> int:
    """Walk AST and accumulate DTD score."""
    total = 0
    is_list_or_set = node.type in ("list", "set")
    for child in node.children:
        if child.type == "function_definition" and child is not top_func:
            continue

        is_element = is_list_or_set and child.type not in _PY_DTD_PUNCTUATION
        if is_element:
            total += depth

        if child.type in _PY_DTD_CONTAINERS:
            total += _dtd_walk_py(child, depth + 1, top_func)
            continue

        if _is_map_call_py(child):
            total += _dtd_walk_py(child, depth + 1, top_func)
            continue

        if child.type in _PY_DTD_PAYLOADS:
            if not is_element:
                total += depth
            total += _dtd_walk_py(child, depth, top_func)
            continue

        total += _dtd_walk_py(child, depth, top_func)

    return total
