"""Shared Halstead collection logic for all languages."""

from __future__ import annotations

from dataclasses import dataclass
from functools import partial
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from tree_sitter import Node

from vibelexity.common import compute_halstead
from vibelexity.models import HalsteadMetrics


@dataclass
class HalsteadConfig:
    comment_types: set[str]
    keyword_operators: set[str]
    operand_types: set[str]
    operator_tokens: set[str]


def _collect_halstead_generic(
    node: Node,
    operators: list[str],
    operands: list[str],
    config: HalsteadConfig,
) -> None:
    """Recursively classify AST nodes as operators or operands.

    Language-specific behavior is controlled entirely by the HalsteadConfig.
    """
    if node.type in config.comment_types:
        return

    if node.type in config.keyword_operators:
        operators.append(node.type)
        for child in node.children:
            _collect_halstead_generic(
                child,
                operators,
                operands,
                config,
            )
        return

    if node.child_count == 0:
        text = node.text.decode()
        if node.type in config.operand_types:
            operands.append(text)
        elif text in config.operator_tokens or node.type in config.operator_tokens:
            operators.append(text)
        return

    for child in node.children:
        _collect_halstead_generic(
            child,
            operators,
            operands,
            config,
        )


def compute_halstead_generic(
    root: Node,
    config: HalsteadConfig,
) -> HalsteadMetrics:
    """Compute Halstead metrics for an AST using the given language config."""
    operators: list[str] = []
    operands: list[str] = []
    _collect_halstead_generic(root, operators, operands, config)
    return compute_halstead(operators, operands)


def make_compute_halstead(
    config: HalsteadConfig,
):
    return partial(
        compute_halstead_generic,
        config=config,
    )
