"""Pure AST utility functions for context enrichment.

Functions that operate only on ``ast`` nodes with no internal
dependencies.  Used by both ``_class_analysis`` and ``enrichment``.
"""

from __future__ import annotations

import ast


class ImportCollector(ast.NodeVisitor):
    """Collect ``from X import Y`` mappings."""

    def __init__(self) -> None:
        """Initialize with an empty name-to-module mapping."""
        self.imported_names: dict[str, str] = {}

    def visit_ImportFrom(
        self,
        node: ast.ImportFrom,
    ) -> None:
        """Record each ``from X import Y`` binding."""
        if node.module:
            for alias in node.names:
                if alias.name != "*":
                    self.imported_names[alias.asname or alias.name] = (
                        node.module
                    )


def bool_literal(node: ast.AST) -> bool | None:
    """Return the boolean value if *node* is a bool constant."""
    if isinstance(node, ast.Constant) and isinstance(node.value, bool):
        return node.value
    return None


def get_expr_name(node: ast.AST | None) -> str | None:
    """Return dotted name for a Name/Attribute chain, or *None*."""
    if node is None:
        return None

    parts: list[str] = []
    current = node
    while True:
        if isinstance(current, ast.Attribute):
            parts.append(current.attr)
            current = current.value
            continue
        if isinstance(current, ast.Call):
            current = current.func
            continue
        if isinstance(current, ast.Name):
            base_name: str | None = current.id
        else:
            base_name = None
        break

    if not parts:
        return base_name

    parts.reverse()
    if base_name is not None:
        parts.insert(0, base_name)
    return ".".join(parts)


def get_node_source(
    node: ast.AST | None,
    module_source: str,
    fallback: str = "...",
) -> str:
    """Extract source text of *node*, falling back to ``ast.unparse``."""
    if node is None:
        return fallback
    source_segment = ast.get_source_segment(module_source, node)
    if source_segment is not None:
        return source_segment
    try:
        return ast.unparse(node)
    except Exception:  # noqa: BLE001
        return fallback


def collect_import_aliases(
    module_tree: ast.Module,
) -> dict[str, str]:
    """Map local import name -> fully-qualified name."""
    aliases: dict[str, str] = {}
    for node in module_tree.body:
        if isinstance(node, ast.Import):
            for alias in node.names:
                bound = alias.asname or alias.name.split(".")[0]
                aliases[bound] = alias.name
        elif isinstance(node, ast.ImportFrom) and node.module:
            for alias in node.names:
                bound = alias.asname or alias.name
                aliases[bound] = f"{node.module}.{alias.name}"
    return aliases


def find_class_node_by_name(
    class_name: str,
    module_tree: ast.Module,
) -> ast.ClassDef | None:
    """Find a ``ClassDef`` by *class_name* in *module_tree*."""
    stack: list[ast.AST] = [module_tree]
    while stack:
        node = stack.pop()
        body = getattr(node, "body", None)
        if body:
            for item in body:
                if isinstance(item, ast.ClassDef):
                    if item.name == class_name:
                        return item
                    stack.append(item)
                elif isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    stack.append(item)
    return None


def collect_existing_class_names(tree: ast.Module) -> set[str]:
    """Return all class names defined in *tree*."""
    class_names: set[str] = set()
    stack: list[ast.AST] = [tree]

    while stack:
        node = stack.pop()
        if isinstance(node, ast.ClassDef):
            class_names.add(node.name)
        if hasattr(node, "body"):
            stack.extend(node.body)
            if hasattr(node, "orelse"):
                stack.extend(node.orelse)
            if hasattr(node, "finalbody"):
                stack.extend(node.finalbody)
            if hasattr(node, "handlers"):
                stack.extend(node.handlers)

    return class_names


def collect_type_names_from_annotation(
    node: ast.expr | None,
) -> set[str]:
    """Recursively collect type names from an annotation node."""
    if isinstance(node, ast.Name):
        return {node.id}
    if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
        return {node.value.id}
    if isinstance(node, ast.Subscript):
        names = collect_type_names_from_annotation(node.value)
        names |= collect_type_names_from_annotation(node.slice)
        return names
    if isinstance(node, ast.BinOp):
        return collect_type_names_from_annotation(
            node.left
        ) | collect_type_names_from_annotation(node.right)
    if isinstance(node, ast.Tuple):
        names = set[str]()
        for elt in node.elts:
            names |= collect_type_names_from_annotation(elt)
        return names
    return set()


def expr_matches_name(
    node: ast.AST | None,
    import_aliases: dict[str, str],
    suffix: str,
) -> bool:
    """Check whether *node*'s resolved name ends with *suffix*."""
    name = get_expr_name(node)
    if name is None:
        return False
    suffix_dot = "." + suffix
    if name == suffix or name.endswith(suffix_dot):
        return True
    resolved_name = import_aliases.get(name)
    return resolved_name is not None and (
        resolved_name == suffix or resolved_name.endswith(suffix_dot)
    )


def resolve_decorator_name(
    expr_name: str,
    import_aliases: dict[str, str],
) -> str:
    """Resolve a decorator expression name via import aliases."""
    resolved = import_aliases.get(expr_name)
    if resolved is not None:
        return resolved
    first_part, sep, rest = expr_name.partition(".")
    if sep:
        root_resolved = import_aliases.get(first_part)
        if root_resolved is not None:
            return f"{root_resolved}.{rest}"
    return expr_name
