"""Cognitive complexity for Python (SonarSource spec)."""

from __future__ import annotations

import functools
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from tree_sitter import Node

from vibelexity.common import _score_boolean_chain as _score_boolean_chain_generic
from vibelexity.metrics.cognitive.common import make_compute_cognitive


# Control flow nodes that add 1 + nesting and then recurse at nesting + 1
_PY_NESTING_INCREMENTORS = {
    "for_statement",
    "while_statement",
    "except_clause",
    "conditional_expression",
    "match_statement",
}


def _flatten_boolean_ops(
    node: Node, ops: list[str], right_subtrees: list[Node]
) -> None:
    """Flatten a left-associative boolean_operator tree into a list of operator strings."""
    if node.type != "boolean_operator":
        return
    children = node.children
    if len(children) >= 3:
        left = children[0]
        operator = children[1]
        right = children[2]
        if left.type == "boolean_operator":
            _flatten_boolean_ops(left, ops, right_subtrees)
        ops.append(operator.text.decode())
        if right.type == "boolean_operator":
            right_subtrees.append(right)


_score_boolean_chain = functools.partial(
    _score_boolean_chain_generic, flatten_fn=_flatten_boolean_ops
)


def _walk(node: Node, nesting: int, func_name: str, top_func: Node) -> int:
    """Recursively walk AST and compute cognitive complexity."""
    total = 0

    for child in node.children:
        if child.type == "function_definition" and child is not top_func:
            continue

        if child.type == "lambda":
            total += _walk(child, nesting + 1, func_name, top_func)
            continue

        if child.type == "if_statement":
            total += 1 + nesting
            total += _walk_if(child, nesting, func_name, top_func)
            continue

        if child.type in _PY_NESTING_INCREMENTORS:
            total += 1 + nesting
            total += _walk(child, nesting + 1, func_name, top_func)
            continue

        if child.type == "boolean_operator":
            total += _score_boolean_chain(child)
            total += _walk_non_boolean(child, nesting, func_name, top_func)
            continue

        if child.type == "call":
            total += _check_recursion(child, func_name)

        total += _walk(child, nesting, func_name, top_func)

    return total


def _walk_if(node: Node, nesting: int, func_name: str, top_func: Node) -> int:
    """Walk children of an if_statement/elif_clause."""
    total = 0
    for child in node.children:
        if child.type == "elif_clause":
            total += 1
            total += _walk_if(child, nesting, func_name, top_func)
        elif child.type == "else_clause":
            total += 1
            total += _walk(child, nesting + 1, func_name, top_func)
        elif child.type == "block":
            total += _walk(child, nesting + 1, func_name, top_func)
        elif child.type == "boolean_operator":
            total += _score_boolean_chain(child)
            total += _walk_non_boolean(child, nesting, func_name, top_func)
        elif child.type == "call":
            total += _check_recursion(child, func_name)
            total += _walk(child, nesting, func_name, top_func)
        else:
            total += _walk(child, nesting, func_name, top_func)
    return total


def _walk_non_boolean(node: Node, nesting: int, func_name: str, top_func: Node) -> int:
    """Walk inside boolean expressions looking only for nested structures and recursion."""
    total = 0
    for child in node.children:
        if child.type == "boolean_operator":
            total += _walk_non_boolean(child, nesting, func_name, top_func)
            continue
        if child.type == "conditional_expression":
            total += 1 + nesting
            total += _walk(child, nesting + 1, func_name, top_func)
            continue
        if child.type == "call":
            total += _check_recursion(child, func_name)
            total += _walk(child, nesting, func_name, top_func)
            continue
        total += _walk_non_boolean(child, nesting, func_name, top_func)
    return total


def _check_recursion(call_node: Node, func_name: str) -> int:
    """Check if a call node is a recursive call to the enclosing function."""
    if not func_name:
        return 0
    for child in call_node.children:
        if child.type == "identifier" and child.text.decode() == func_name:
            return 1
        if child.type == "argument_list":
            break
    return 0


compute_cognitive_py = make_compute_cognitive(_walk)
