#exonware/xwquery/src/exonware/xwquery/query/adapters/ast_utils.py
"""
AST Utility Functions for xwquery
Provides helper functions for traversing and extracting data from AST nodes
generated by the xwsystem.syntax Grammar engine.
These utilities support converting AST nodes from various query language grammars
into xwquery QueryAction trees.
Company: eXonware.com
Author: eXonware Backend Team
Email: connect@exonware.com
Version: 0.9.0.5
Generation Date: 29-Oct-2024
Priority Alignment:
- Security (#1): Safe AST traversal with depth limits
- Usability (#2): Clean, intuitive API for AST operations
- Maintainability (#3): Reusable functions, DRY principle
- Performance (#4): Efficient tree traversal algorithms
- Extensibility (#5): Easy to extend for new grammar formats
"""

from typing import Any
from exonware.xwsyntax import ParseNode
# Maximum depth for AST traversal (prevent infinite recursion)
from collections.abc import Callable, Iterator
MAX_AST_DEPTH = 1000


def find_node_by_type(ast: ParseNode, node_type: str) -> ParseNode | None:
    """
    Find first node of specified type in AST tree (depth-first search).
    Args:
        ast: Root AST node to search
        node_type: Type of node to find (e.g., 'select_statement', 'where_clause')
    Returns:
        First matching ParseNode, or None if not found
    Security: Depth-limited to prevent DoS attacks with deeply nested queries
    Example:
        >>> select_node = find_node_by_type(ast, 'select_statement')
        >>> where_node = find_node_by_type(ast, 'where_clause')
    """
    if not ast:
        return None
    # Check current node
    if ast.type == node_type:
        return ast
    # Search children
    if hasattr(ast, 'children') and ast.children:
        for child in ast.children:
            result = find_node_by_type(child, node_type)
            if result:
                return result
    return None


def find_all_nodes_by_type(ast: ParseNode, node_type: str, max_depth: int = MAX_AST_DEPTH) -> list[ParseNode]:
    """
    Find all nodes of specified type in AST tree (depth-first search).
    Args:
        ast: Root AST node to search
        node_type: Type of node to find
        max_depth: Maximum traversal depth (security limit)
    Returns:
        List of all matching ParseNodes (may be empty)
    Security: Depth-limited to prevent resource exhaustion
    Example:
        >>> column_nodes = find_all_nodes_by_type(ast, 'column_ref')
        >>> for col in column_nodes:
        ...     print(col.value)
    """
    results = []
    def _search(node: ParseNode, depth: int):
        if depth > max_depth:
            return  # Stop at max depth
        if not node:
            return
        # Check current node
        if node.type == node_type:
            results.append(node)
        # Search children
        if hasattr(node, 'children') and node.children:
            for child in node.children:
                _search(child, depth + 1)
    _search(ast, 0)
    return results


def extract_node_value(ast: ParseNode) -> Any:
    """
    Extract value from AST node.
    Handles both terminal nodes (with direct values) and non-terminal nodes
    (which may have values in children).
    Args:
        ast: ParseNode to extract value from
    Returns:
        Extracted value (str, int, float, bool, None, or List)
    Example:
        >>> value = extract_node_value(identifier_node)  # Returns "users"
        >>> values = extract_node_value(list_node)  # Returns ["col1", "col2"]
    """
    if not ast:
        return None
    # Check for direct value
    if hasattr(ast, 'value') and ast.value is not None:
        return ast.value
    # For terminal nodes (type='terminal')
    if ast.type == 'terminal':
        return ast.value if hasattr(ast, 'value') else None
    # For nodes with single child, extract child value
    if hasattr(ast, 'children') and ast.children:
        if len(ast.children) == 1:
            return extract_node_value(ast.children[0])
        # For multiple children, extract all values
        values = []
        for child in ast.children:
            val = extract_node_value(child)
            if val is not None:
                values.append(val)
        return values if values else None
    return None


def traverse_depth_first(ast: ParseNode, max_depth: int = MAX_AST_DEPTH) -> Iterator[ParseNode]:
    """
    Traverse AST tree depth-first (pre-order).
    Args:
        ast: Root AST node
        max_depth: Maximum traversal depth (security limit)
    Yields:
        ParseNode in depth-first pre-order
    Security: Depth-limited to prevent stack overflow
    Example:
        >>> for node in traverse_depth_first(ast):
        ...     print(node.type, node.value)
    """
    def _traverse(node: ParseNode, depth: int):
        if depth > max_depth:
            return  # Stop at max depth
        if not node:
            return
        # Yield current node (pre-order)
        yield node
        # Traverse children
        if hasattr(node, 'children') and node.children:
            for child in node.children:
                yield from _traverse(child, depth + 1)
    yield from _traverse(ast, 0)


def traverse_breadth_first(ast: ParseNode) -> Iterator[ParseNode]:
    """
    Traverse AST tree breadth-first (level-order).
    Args:
        ast: Root AST node
    Yields:
        ParseNode in breadth-first order
    Example:
        >>> for node in traverse_breadth_first(ast):
        ...     print(f"Level: {node.metadata.get('depth', 0)}, Type: {node.type}")
    """
    if not ast:
        return
    queue = [ast]
    while queue:
        node = queue.pop(0)
        yield node
        if hasattr(node, 'children') and node.children:
            queue.extend(node.children)


def find_parent(ast: ParseNode, target_node: ParseNode) -> ParseNode | None:
    """
    Find parent of target node in AST tree.
    Args:
        ast: Root AST node
        target_node: Node to find parent of
    Returns:
        Parent ParseNode, or None if target is root or not found
    Example:
        >>> parent = find_parent(ast, where_node)
    """
    if not ast or not target_node:
        return None
    # Check if any child matches target
    if hasattr(ast, 'children') and ast.children:
        for child in ast.children:
            if child is target_node:
                return ast
            # Recurse into children
            parent = find_parent(child, target_node)
            if parent:
                return parent
    return None


def get_node_text(ast: ParseNode) -> str:
    """
    Get text representation of node and all descendants.
    Concatenates all terminal values in the subtree.
    Args:
        ast: Root AST node
    Returns:
        Text representation of subtree
    Example:
        >>> text = get_node_text(expression_node)  # "age > 18"
    """
    if not ast:
        return ""
    # Terminal node - return value directly
    if ast.type == 'terminal' or not hasattr(ast, 'children'):
        return str(ast.value) if ast.value is not None else ""
    # Non-terminal - concatenate children
    parts = []
    for child in ast.children:
        text = get_node_text(child)
        if text:
            parts.append(text)
    return " ".join(parts)


def find_nodes_by_predicate(
    ast: ParseNode,
    predicate: Callable[[ParseNode], bool],
    max_depth: int = MAX_AST_DEPTH
) -> list[ParseNode]:
    """
    Find all nodes matching a predicate function.
    Args:
        ast: Root AST node
        predicate: Function that returns True for matching nodes
        max_depth: Maximum traversal depth
    Returns:
        List of matching nodes
    Example:
        >>> # Find all identifier nodes
        >>> identifiers = find_nodes_by_predicate(
        ...     ast,
        ...     lambda n: n.type == 'identifier'
        ... )
        >>> 
        >>> # Find all nodes with specific value
        >>> age_nodes = find_nodes_by_predicate(
        ...     ast,
        ...     lambda n: n.value == 'age'
        ... )
    """
    results = []
    def _search(node: ParseNode, depth: int):
        if depth > max_depth:
            return
        if not node:
            return
        # Check predicate
        if predicate(node):
            results.append(node)
        # Search children
        if hasattr(node, 'children') and node.children:
            for child in node.children:
                _search(child, depth + 1)
    _search(ast, 0)
    return results


def extract_keyword_node(ast: ParseNode, keyword: str) -> ParseNode | None:
    """
    Find node following a specific keyword.
    Useful for extracting clauses like WHERE, FROM, ORDER BY, etc.
    Args:
        ast: Root AST node
        keyword: Keyword to search for (case-insensitive)
    Returns:
        Node following the keyword, or None if not found
    Example:
        >>> where_node = extract_keyword_node(ast, 'WHERE')
        >>> from_node = extract_keyword_node(ast, 'FROM')
    """
    keyword_upper = keyword.upper()
    for i, node in enumerate(traverse_depth_first(ast)):
        # Check if node matches keyword
        if node.value and str(node.value).upper() == keyword_upper:
            # Find parent to get next sibling
            parent = find_parent(ast, node)
            if parent and hasattr(parent, 'children'):
                # Find node's index
                try:
                    node_idx = parent.children.index(node)
                    # Return next sibling if exists
                    if node_idx + 1 < len(parent.children):
                        return parent.children[node_idx + 1]
                except (ValueError, IndexError):
                    pass
    return None


def ast_to_dict(ast: ParseNode, max_depth: int = 10) -> dict[str, Any]:
    """
    Convert AST to dictionary representation (for debugging/logging).
    Args:
        ast: Root AST node
        max_depth: Maximum depth to convert
    Returns:
        Dictionary representation of AST
    Example:
        >>> ast_dict = ast_to_dict(ast)
        >>> from exonware.xwsystem import get_serializer, JsonSerializer
        >>> js = get_serializer(JsonSerializer)
        >>> out = js.dumps(ast_dict); print(out.decode('utf-8') if isinstance(out, bytes) else out)
    """
    if not ast or max_depth <= 0:
        return {}
    result = {
        "type": ast.type,
        "value": ast.value if hasattr(ast, 'value') else None,
        "metadata": ast.metadata if hasattr(ast, 'metadata') else {},
    }
    if hasattr(ast, 'children') and ast.children:
        result["children"] = [
            ast_to_dict(child, max_depth - 1) for child in ast.children
        ]
    return result


def count_nodes(ast: ParseNode) -> int:
    """
    Count total number of nodes in AST tree.
    Args:
        ast: Root AST node
    Returns:
        Total node count
    Example:
        >>> node_count = count_nodes(ast)
        >>> print(f"AST has {node_count} nodes")
    """
    if not ast:
        return 0
    count = 1  # Count current node
    if hasattr(ast, 'children') and ast.children:
        for child in ast.children:
            count += count_nodes(child)
    return count


def get_ast_depth(ast: ParseNode) -> int:
    """
    Get maximum depth of AST tree.
    Args:
        ast: Root AST node
    Returns:
        Maximum depth (root = 0)
    Example:
        >>> depth = get_ast_depth(ast)
        >>> print(f"AST depth: {depth}")
    """
    if not ast:
        return 0
    if not hasattr(ast, 'children') or not ast.children:
        return 0
    max_child_depth = 0
    for child in ast.children:
        child_depth = get_ast_depth(child)
        max_child_depth = max(max_child_depth, child_depth)
    return 1 + max_child_depth


def collect_terminals(ast: ParseNode) -> list[ParseNode]:
    """
    Collect all terminal (leaf) nodes from AST tree.
    Args:
        ast: Root AST node
    Returns:
        List of terminal nodes
    Example:
        >>> terminals = collect_terminals(ast)
        >>> text = " ".join(t.value for t in terminals if t.value)
    """
    terminals = []
    for node in traverse_depth_first(ast):
        # Terminal if no children or type is 'terminal'
        if node.type == 'terminal' or not hasattr(node, 'children') or not node.children:
            terminals.append(node)
    return terminals


def extract_text_from_subtree(ast: ParseNode, separator: str = " ") -> str:
    """
    Extract all text from subtree (concatenate terminal values).
    Args:
        ast: Root AST node
        separator: Separator between values
    Returns:
        Concatenated text from all terminals
    Example:
        >>> text = extract_text_from_subtree(expression_node)
        >>> print(text)  # "age > 18"
    """
    terminals = collect_terminals(ast)
    texts = [str(t.value) for t in terminals if t.value is not None]
    return separator.join(texts)
# Convenience functions for common patterns

def has_child_of_type(ast: ParseNode, node_type: str) -> bool:
    """Check if node has a direct child of specified type."""
    if not ast or not hasattr(ast, 'children'):
        return False
    return any(child.type == node_type for child in ast.children)


def get_child_by_type(ast: ParseNode, node_type: str) -> ParseNode | None:
    """Get first direct child of specified type."""
    if not ast or not hasattr(ast, 'children'):
        return None
    for child in ast.children:
        if child.type == node_type:
            return child
    return None


def get_children_by_type(ast: ParseNode, node_type: str) -> list[ParseNode]:
    """Get all direct children of specified type."""
    if not ast or not hasattr(ast, 'children'):
        return []
    return [child for child in ast.children if child.type == node_type]


def is_terminal(node: ParseNode) -> bool:
    """Check if node is a terminal (leaf) node."""
    return node.type == 'terminal' or not hasattr(node, 'children') or not node.children


def get_node_at_path(ast: ParseNode, path: list[str]) -> ParseNode | None:
    """
    Navigate AST using a path of node types.
    Args:
        ast: Root AST node
        path: List of node types to traverse (e.g., ['query', 'select_stmt', 'where_clause'])
    Returns:
        Node at path, or None if path doesn't exist
    Example:
        >>> where = get_node_at_path(ast, ['query', 'select_statement', 'where_clause'])
    """
    current = ast
    for node_type in path:
        if not current:
            return None
        # Find child of specified type
        current = get_child_by_type(current, node_type)
    return current


def extract_identifier_chain(ast: ParseNode) -> list[str]:
    """
    Extract chain of identifiers (for property access like a.b.c).
    Args:
        ast: ParseNode representing property reference
    Returns:
        List of identifier strings
    Example:
        >>> chain = extract_identifier_chain(property_ref_node)
        >>> print(".".join(chain))  # "users.name"
    """
    identifiers = []
    for node in traverse_depth_first(ast):
        if node.type.lower() in ('identifier', 'ncname', 'name'):
            value = extract_node_value(node)
            if value and isinstance(value, str):
                identifiers.append(value)
    return identifiers


def print_ast_tree(ast: ParseNode, indent: int = 0, max_depth: int = 10):
    """
    Print AST tree structure (for debugging).
    Args:
        ast: Root AST node
        indent: Current indentation level
        max_depth: Maximum depth to print
    Example:
        >>> print_ast_tree(ast)
        query
          select_statement
            select_list
              column_ref: users
            from_clause
              table_ref: users
    """
    if not ast or indent > max_depth:
        return
    prefix = "  " * indent
    value_str = f": {ast.value}" if ast.value is not None else ""
    print(f"{prefix}{ast.type}{value_str}")
    if hasattr(ast, 'children') and ast.children:
        for child in ast.children:
            print_ast_tree(child, indent + 1, max_depth)
