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

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from tree_sitter import Node

from vibelexity.parsers.go import is_named_func_literal

_GO_DTD_MAP_NAMES = {"Map", "Filter", "Reduce"}


def compute_dtd_go(func_node: Node) -> int:
    """Compute Data Transformation Density for a Go function node."""
    body = func_node.child_by_field_name("body")
    if body and body.type == "block":
        return _dtd_walk_go(body, depth=0, top_func=func_node)
    for child in func_node.children:
        if child.type == "block":
            return _dtd_walk_go(child, depth=0, top_func=func_node)
    return 0


def _is_map_call_go(node: Node) -> bool:
    """Check if a call_expression is a Map/Filter/Reduce call."""
    if node.type != "call_expression":
        return False
    func = node.child_by_field_name("function")
    if func and func.type == "selector_expression":
        field = func.child_by_field_name("field")
        if field and field.text.decode() in _GO_DTD_MAP_NAMES:
            return True
    return False


_GO_DTD_NESTED_FUNCS = {"function_declaration", "method_declaration"}
_GO_DTD_LITERAL_ELEMENTS = {"keyed_element", "literal_element"}


def _dtd_walk_go(node: Node, depth: int, top_func: Node) -> int:
    """Walk AST and accumulate DTD score."""
    total = 0
    is_literal = node.type == "literal_value"
    for child in node.children:
        if child.type in _GO_DTD_NESTED_FUNCS and child is not top_func:
            continue
        if child.type == "func_literal" and child is not top_func:
            if not is_named_func_literal(child):
                total += _dtd_walk_go(child, depth, top_func)
            continue

        if child.type == "literal_value" or _is_map_call_go(child):
            total += _dtd_walk_go(child, depth + 1, top_func)
            continue

        if is_literal and child.type in _GO_DTD_LITERAL_ELEMENTS:
            total += depth

        total += _dtd_walk_go(child, depth, top_func)

    return total
