from typing import Literal, TypedDict, type_check_only
from typing_extensions import NotRequired, TypeAlias

_NodeType: TypeAlias = Literal[
    "comparator",
    "current",
    "expref",
    "function_expression",
    "field",
    "filter_projection",
    "flatten",
    "identity",
    "index",
    "index_expression",
    "key_val_pair",
    "literal",
    "multi_select_dict",
    "multi_select_list",
    "or_expression",
    "and_expression",
    "not_expression",
    "pipe",
    "projection",
    "subexpression",
    "slice",
    "value_projection",
]

@type_check_only
class _ASTNode(TypedDict):
    type: _NodeType
    children: list[_ASTNode]
    value: NotRequired[str]

def comparator(name: str, first: _ASTNode, second: _ASTNode) -> _ASTNode: ...
def current_node() -> _ASTNode: ...
def expref(expression: _ASTNode) -> _ASTNode: ...
def function_expression(name: str, args: list[_ASTNode]) -> _ASTNode: ...
def field(name: str) -> _ASTNode: ...
def filter_projection(left: _ASTNode, right: _ASTNode, comparator: _ASTNode) -> _ASTNode: ...
def flatten(node: _ASTNode) -> _ASTNode: ...
def identity() -> _ASTNode: ...
def index(index: str) -> _ASTNode: ...
def index_expression(children: list[_ASTNode]) -> _ASTNode: ...
def key_val_pair(key_name: str, node: _ASTNode) -> _ASTNode: ...
def literal(literal_value: str) -> _ASTNode: ...
def multi_select_dict(nodes: list[_ASTNode]) -> _ASTNode: ...
def multi_select_list(nodes: list[_ASTNode]) -> _ASTNode: ...
def or_expression(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
def and_expression(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
def not_expression(expr: _ASTNode) -> _ASTNode: ...
def pipe(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
def projection(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
def subexpression(children: list[_ASTNode]) -> _ASTNode: ...
def slice(start: _ASTNode, end: _ASTNode, step: _ASTNode) -> _ASTNode: ...
def value_projection(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
