"""
blib2to3 Node/Leaf transformation-related utility functions.
"""

from collections.abc import Iterator
from typing import Final, Generic, Literal, TypeGuard, TypeVar, Union

from mypy_extensions import mypyc_attr

from black.cache import CACHE_DIR
from black.mode import Mode, Preview
from black.strings import get_string_prefix, has_triple_quotes
from blib2to3 import pygram
from blib2to3.pgen2 import token
from blib2to3.pytree import NL, Leaf, Node, type_repr

pygram.initialize(CACHE_DIR)
syms: Final = pygram.python_symbols


# types
T = TypeVar("T")
LN = Union[Leaf, Node]
LeafID = int
NodeType = int


WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
STATEMENT: Final = {
    syms.if_stmt,
    syms.while_stmt,
    syms.for_stmt,
    syms.try_stmt,
    syms.except_clause,
    syms.with_stmt,
    syms.funcdef,
    syms.classdef,
    syms.match_stmt,
    syms.case_block,
}
STANDALONE_COMMENT: Final = 153
token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
LOGIC_OPERATORS: Final = {"and", "or"}
COMPARATORS: Final = {
    token.LESS,
    token.GREATER,
    token.EQEQUAL,
    token.NOTEQUAL,
    token.LESSEQUAL,
    token.GREATEREQUAL,
}
MATH_OPERATORS: Final = {
    token.VBAR,
    token.CIRCUMFLEX,
    token.AMPER,
    token.LEFTSHIFT,
    token.RIGHTSHIFT,
    token.PLUS,
    token.MINUS,
    token.STAR,
    token.SLASH,
    token.DOUBLESLASH,
    token.PERCENT,
    token.AT,
    token.TILDE,
    token.DOUBLESTAR,
}
STARS: Final = {token.STAR, token.DOUBLESTAR}
VARARGS_SPECIALS: Final = STARS | {token.SLASH}
VARARGS_PARENTS: Final = {
    syms.arglist,
    syms.argument,  # double star in arglist
    syms.trailer,  # single argument to call
    syms.typedargslist,
    syms.varargslist,  # lambdas
}
UNPACKING_PARENTS: Final = {
    syms.atom,  # single element of a list or set literal
    syms.dictsetmaker,
    syms.listmaker,
    syms.testlist_gexp,
    syms.testlist_star_expr,
    syms.subject_expr,
    syms.pattern,
}
TEST_DESCENDANTS: Final = {
    syms.test,
    syms.lambdef,
    syms.or_test,
    syms.and_test,
    syms.not_test,
    syms.comparison,
    syms.star_expr,
    syms.expr,
    syms.xor_expr,
    syms.and_expr,
    syms.shift_expr,
    syms.arith_expr,
    syms.trailer,
    syms.term,
    syms.power,
    syms.namedexpr_test,
}
TYPED_NAMES: Final = {syms.tname, syms.tname_star}
ASSIGNMENTS: Final = {
    "=",
    "+=",
    "-=",
    "*=",
    "@=",
    "/=",
    "%=",
    "&=",
    "|=",
    "^=",
    "<<=",
    ">>=",
    "**=",
    "//=",
    ":",
}

IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
BRACKET: Final = {
    token.LPAR: token.RPAR,
    token.LSQB: token.RSQB,
    token.LBRACE: token.RBRACE,
}
OPENING_BRACKETS: Final = set(BRACKET.keys())
CLOSING_BRACKETS: Final = set(BRACKET.values())
BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS
ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {
    token.COMMA,
    STANDALONE_COMMENT,
    token.FSTRING_MIDDLE,
    token.FSTRING_END,
    token.TSTRING_MIDDLE,
    token.TSTRING_END,
    token.BANG,
}

RARROW = 55


@mypyc_attr(allow_interpreted_subclasses=True)
class Visitor(Generic[T]):
    """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""

    def visit(self, node: LN) -> Iterator[T]:
        """Main method to visit `node` and its children.

        It tries to find a `visit_*()` method for the given `node.type`, like
        `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
        If no dedicated `visit_*()` method is found, chooses `visit_default()`
        instead.

        Then yields objects of type `T` from the selected visitor.
        """
        if node.type < 256:
            name = token.tok_name[node.type]
        else:
            name = str(type_repr(node.type))
        # We explicitly branch on whether a visitor exists (instead of
        # using self.visit_default as the default arg to getattr) in order
        # to save needing to create a bound method object and so mypyc can
        # generate a native call to visit_default.
        visitf = getattr(self, f"visit_{name}", None)
        if visitf:
            yield from visitf(node)
        else:
            yield from self.visit_default(node)

    def visit_default(self, node: LN) -> Iterator[T]:
        """Default `visit_*()` implementation. Recurses to children of `node`."""
        if isinstance(node, Node):
            for child in node.children:
                yield from self.visit(child)


def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str:
    """Return whitespace prefix if needed for the given `leaf`.

    `complex_subscript` signals whether the given leaf is part of a subscription
    which has non-trivial arguments, like arithmetic expressions or function calls.
    """
    NO: Final[str] = ""
    SPACE: Final[str] = " "
    DOUBLESPACE: Final[str] = "  "
    t = leaf.type
    p = leaf.parent
    v = leaf.value
    if t in ALWAYS_NO_SPACE:
        return NO

    if t == token.COMMENT:
        return DOUBLESPACE

    assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
    if t == token.COLON and p.type not in {
        syms.subscript,
        syms.subscriptlist,
        syms.sliceop,
    }:
        return NO

    if t == token.LBRACE and p.type in (
        syms.fstring_replacement_field,
        syms.tstring_replacement_field,
    ):
        return NO

    prev = leaf.prev_sibling
    if not prev:
        prevp = preceding_leaf(p)
        if not prevp or prevp.type in OPENING_BRACKETS:
            return NO

        if t == token.COLON:
            if prevp.type == token.COLON:
                return NO

            elif prevp.type != token.COMMA and not complex_subscript:
                return NO

            return SPACE

        if prevp.type == token.EQUAL:
            if prevp.parent:
                if prevp.parent.type in {
                    syms.arglist,
                    syms.argument,
                    syms.parameters,
                    syms.varargslist,
                }:
                    return NO

                elif prevp.parent.type == syms.typedargslist:
                    # A bit hacky: if the equal sign has whitespace, it means we
                    # previously found it's a typed argument.  So, we're using
                    # that, too.
                    return prevp.prefix

        elif (
            prevp.type == token.STAR
            and parent_type(prevp) == syms.star_expr
            and parent_type(prevp.parent) in (syms.subscriptlist, syms.tname_star)
        ):
            # No space between typevar tuples or unpacking them.
            return NO

        elif prevp.type in VARARGS_SPECIALS:
            if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
                return NO

        elif prevp.type == token.COLON:
            if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
                return SPACE if complex_subscript else NO

        elif (
            prevp.parent
            and prevp.parent.type == syms.factor
            and prevp.type in MATH_OPERATORS
        ):
            return NO

        elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
            # no space in decorators
            return NO

    elif prev.type in OPENING_BRACKETS:
        return NO

    elif prev.type == token.BANG:
        return NO

    if p.type in {syms.parameters, syms.arglist}:
        # untyped function signatures or calls
        if not prev or prev.type != token.COMMA:
            return NO

    elif p.type == syms.varargslist:
        # lambdas
        if prev and prev.type != token.COMMA:
            return NO

    elif p.type == syms.typedargslist:
        # typed function signatures
        if not prev:
            return NO

        if t == token.EQUAL:
            if prev.type not in TYPED_NAMES:
                return NO

        elif prev.type == token.EQUAL:
            # A bit hacky: if the equal sign has whitespace, it means we
            # previously found it's a typed argument.  So, we're using that, too.
            return prev.prefix

        elif prev.type != token.COMMA:
            return NO

    elif p.type in TYPED_NAMES:
        # type names
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type != token.COMMA:
                return NO

    elif p.type == syms.trailer:
        # attributes and calls
        if t == token.LPAR or t == token.RPAR:
            return NO

        if not prev:
            if t == token.DOT or t == token.LSQB:
                return NO

        elif prev.type != token.COMMA:
            return NO

    elif p.type == syms.argument:
        # single argument
        if t == token.EQUAL:
            return NO

        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type == token.LPAR:
                return NO

        elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
            return NO

    elif p.type == syms.decorator:
        # decorators
        return NO

    elif p.type == syms.dotted_name:
        if prev:
            return NO

        prevp = preceding_leaf(p)
        if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
            return NO

    elif p.type == syms.classdef:
        if t == token.LPAR:
            return NO

        if prev and prev.type == token.LPAR:
            return NO

    elif p.type in {syms.subscript, syms.sliceop}:
        # indexing
        if not prev:
            assert p.parent is not None, "subscripts are always parented"
            if p.parent.type == syms.subscriptlist:
                return SPACE

            return NO

        elif t == token.COLONEQUAL or prev.type == token.COLONEQUAL:
            return SPACE

        elif not complex_subscript:
            return NO

    elif p.type == syms.atom:
        if prev and t == token.DOT:
            # dots, but not the first one.
            return NO

    elif p.type == syms.dictsetmaker:
        # dict unpacking
        if prev and prev.type == token.DOUBLESTAR:
            return NO

    elif p.type in {syms.factor, syms.star_expr}:
        # unary ops
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type in OPENING_BRACKETS:
                return NO

            prevp_parent = prevp.parent
            assert prevp_parent is not None
            if prevp.type == token.COLON and prevp_parent.type in {
                syms.subscript,
                syms.sliceop,
            }:
                return NO

            elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
                return NO

        elif t in {token.NAME, token.NUMBER, token.STRING}:
            return NO

    elif p.type == syms.import_from:
        if t == token.DOT:
            if prev and prev.type == token.DOT:
                return NO

        elif t == token.NAME:
            if v == "import":
                return SPACE

            if prev and prev.type == token.DOT:
                return NO

    elif p.type == syms.sliceop:
        return NO

    elif p.type == syms.except_clause:
        if t == token.STAR:
            return NO

    if Preview.simplify_power_operator_hugging in mode:
        # Power operator hugging
        if t == token.DOUBLESTAR and is_simple_exponentiation(p):
            return NO
        prevp = preceding_leaf(leaf)
        if prevp and prevp.type == token.DOUBLESTAR:
            if prevp.parent and is_simple_exponentiation(prevp.parent):
                return NO

    return SPACE


def make_simple_prefix(nl_count: int, form_feed: bool, empty_line: str = "\n") -> str:
    """Generate a normalized prefix string."""
    if form_feed:
        return (empty_line * (nl_count - 1)) + "\f" + empty_line
    return empty_line * nl_count


def preceding_leaf(node: LN | None) -> Leaf | None:
    """Return the first leaf that precedes `node`, if any."""
    while node:
        res = node.prev_sibling
        if res:
            if isinstance(res, Leaf):
                return res

            try:
                return list(res.leaves())[-1]

            except IndexError:
                return None

        node = node.parent
    return None


def prev_siblings_are(node: LN | None, tokens: list[NodeType | None]) -> bool:
    """Return if the `node` and its previous siblings match types against the provided
    list of tokens; the provided `node`has its type matched against the last element in
    the list.  `None` can be used as the first element to declare that the start of the
    list is anchored at the start of its parent's children."""
    if not tokens:
        return True
    if tokens[-1] is None:
        return node is None
    if not node:
        return False
    if node.type != tokens[-1]:
        return False
    return prev_siblings_are(node.prev_sibling, tokens[:-1])


def parent_type(node: LN | None) -> NodeType | None:
    """
    Returns:
        @node.parent.type, if @node is not None and has a parent.
            OR
        None, otherwise.
    """
    if node is None or node.parent is None:
        return None

    return node.parent.type


def child_towards(ancestor: Node, descendant: LN) -> LN | None:
    """Return the child of `ancestor` that contains `descendant`."""
    node: LN | None = descendant
    while node and node.parent != ancestor:
        node = node.parent
    return node


def replace_child(old_child: LN, new_child: LN) -> None:
    """
    Side Effects:
        * If @old_child.parent is set, replace @old_child with @new_child in
        @old_child's underlying Node structure.
            OR
        * Otherwise, this function does nothing.
    """
    parent = old_child.parent
    if not parent:
        return

    child_idx = old_child.remove()
    if child_idx is not None:
        parent.insert_child(child_idx, new_child)


def container_of(leaf: Leaf) -> LN:
    """Return `leaf` or one of its ancestors that is the topmost container of it.

    By "container" we mean a node where `leaf` is the very first child.
    """
    same_prefix = leaf.prefix
    container: LN = leaf
    while container:
        parent = container.parent
        if parent is None:
            break

        if parent.children[0].prefix != same_prefix:
            break

        if parent.type == syms.file_input:
            break

        if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
            break

        container = parent
    return container


def first_leaf_of(node: LN) -> Leaf | None:
    """Returns the first leaf of the node tree."""
    if isinstance(node, Leaf):
        return node
    if node.children:
        return first_leaf_of(node.children[0])
    else:
        return None


def is_arith_like(node: LN) -> bool:
    """Whether node is an arithmetic or a binary arithmetic expression"""
    return node.type in {
        syms.arith_expr,
        syms.shift_expr,
        syms.xor_expr,
        syms.and_expr,
    }


def is_simple_exponentiation(node: LN) -> bool:
    """Whether whitespace around `**` should be removed."""

    def is_simple(node: LN) -> bool:
        if isinstance(node, Leaf):
            return node.type in (token.NAME, token.NUMBER, token.DOT, token.DOUBLESTAR)
        elif node.type == syms.factor:  # unary operators
            return is_simple(node.children[1])
        else:
            return all(is_simple(child) for child in node.children)

    return (
        node.type == syms.power
        and len(node.children) >= 3
        and node.children[-2].type == token.DOUBLESTAR
        and is_simple(node)
    )


def is_docstring(node: NL) -> bool:
    if isinstance(node, Leaf):
        if node.type != token.STRING:
            return False

        prefix = get_string_prefix(node.value)
        if set(prefix).intersection("bBfF"):
            return False

    if (
        node.parent
        and node.parent.type == syms.simple_stmt
        and not node.parent.prev_sibling
        and node.parent.parent
        and node.parent.parent.type == syms.file_input
    ):
        return True

    if prev_siblings_are(
        node.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
    ):
        return True

    # Multiline docstring on the same line as the `def`.
    if prev_siblings_are(node.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
        # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
        # grammar. We're safe to return True without further checks.
        return True

    return False


def is_empty_tuple(node: LN) -> bool:
    """Return True if `node` holds an empty tuple."""
    return (
        node.type == syms.atom
        and len(node.children) == 2
        and node.children[0].type == token.LPAR
        and node.children[1].type == token.RPAR
    )


def is_one_tuple(node: LN) -> bool:
    """Return True if `node` holds a tuple with one element, with or without parens."""
    if node.type == syms.atom:
        gexp = unwrap_singleton_parenthesis(node)
        if gexp is None or gexp.type != syms.testlist_gexp:
            return False

        return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA

    return (
        node.type in IMPLICIT_TUPLE
        and len(node.children) == 2
        and node.children[1].type == token.COMMA
    )


def is_tuple(node: LN) -> bool:
    """Return True if `node` holds a tuple."""
    if node.type != syms.atom:
        return False
    gexp = unwrap_singleton_parenthesis(node)
    if gexp is None or gexp.type != syms.testlist_gexp:
        return False

    return True


def is_tuple_containing_walrus(node: LN) -> bool:
    """Return True if `node` holds a tuple that contains a walrus operator."""
    if node.type != syms.atom:
        return False
    gexp = unwrap_singleton_parenthesis(node)
    if gexp is None or gexp.type != syms.testlist_gexp:
        return False

    return any(child.type == syms.namedexpr_test for child in gexp.children)


def is_tuple_containing_star(node: LN) -> bool:
    """Return True if `node` holds a tuple that contains a star operator."""
    if node.type != syms.atom:
        return False
    gexp = unwrap_singleton_parenthesis(node)
    if gexp is None or gexp.type != syms.testlist_gexp:
        return False

    return any(child.type == syms.star_expr for child in gexp.children)


def is_generator(node: LN) -> bool:
    """Return True if `node` holds a generator."""
    if node.type != syms.atom:
        return False
    gexp = unwrap_singleton_parenthesis(node)
    if gexp is None or gexp.type != syms.testlist_gexp:
        return False

    return any(child.type == syms.old_comp_for for child in gexp.children)


def is_one_sequence_between(
    opening: Leaf,
    closing: Leaf,
    leaves: list[Leaf],
    brackets: tuple[int, int] = (token.LPAR, token.RPAR),
) -> bool:
    """Return True if content between `opening` and `closing` is a one-sequence."""
    if (opening.type, closing.type) != brackets:
        return False

    depth = closing.bracket_depth + 1
    for _opening_index, leaf in enumerate(leaves):
        if leaf is opening:
            break

    else:
        return False

    commas = 0
    _opening_index += 1
    for leaf in leaves[_opening_index:]:
        if leaf is closing:
            break

        bracket_depth = leaf.bracket_depth
        if bracket_depth == depth and leaf.type == token.COMMA:
            commas += 1
            if leaf.parent and leaf.parent.type in {
                syms.arglist,
                syms.typedargslist,
            }:
                commas += 1
                break

    return commas < 2


def is_walrus_assignment(node: LN) -> bool:
    """Return True iff `node` is of the shape ( test := test )"""
    inner = unwrap_singleton_parenthesis(node)
    return inner is not None and inner.type == syms.namedexpr_test


def is_simple_decorator_trailer(node: LN, last: bool = False) -> bool:
    """Return True iff `node` is a trailer valid in a simple decorator"""
    return node.type == syms.trailer and (
        (
            len(node.children) == 2
            and node.children[0].type == token.DOT
            and node.children[1].type == token.NAME
        )
        # last trailer can be an argument-less parentheses pair
        or (
            last
            and len(node.children) == 2
            and node.children[0].type == token.LPAR
            and node.children[1].type == token.RPAR
        )
        # last trailer can be arguments
        or (
            last
            and len(node.children) == 3
            and node.children[0].type == token.LPAR
            # and node.children[1].type == syms.argument
            and node.children[2].type == token.RPAR
        )
    )


def is_simple_decorator_expression(node: LN) -> bool:
    """Return True iff `node` could be a 'dotted name' decorator

    This function takes the node of the 'namedexpr_test' of the new decorator
    grammar and test if it would be valid under the old decorator grammar.

    The old grammar was: decorator: @ dotted_name [arguments] NEWLINE
    The new grammar is : decorator: @ namedexpr_test NEWLINE
    """
    if node.type == token.NAME:
        return True
    if node.type == syms.power:
        if node.children:
            return (
                node.children[0].type == token.NAME
                and all(map(is_simple_decorator_trailer, node.children[1:-1]))
                and (
                    len(node.children) < 2
                    or is_simple_decorator_trailer(node.children[-1], last=True)
                )
            )
    return False


def is_yield(node: LN) -> bool:
    """Return True if `node` holds a `yield` or `yield from` expression."""
    if node.type == syms.yield_expr:
        return True

    if is_name_token(node) and node.value == "yield":
        return True

    if node.type != syms.atom:
        return False

    if len(node.children) != 3:
        return False

    lpar, expr, rpar = node.children
    if lpar.type == token.LPAR and rpar.type == token.RPAR:
        return is_yield(expr)

    return False


def is_vararg(leaf: Leaf, within: set[NodeType]) -> bool:
    """Return True if `leaf` is a star or double star in a vararg or kwarg.

    If `within` includes VARARGS_PARENTS, this applies to function signatures.
    If `within` includes UNPACKING_PARENTS, it applies to right hand-side
    extended iterable unpacking (PEP 3132) and additional unpacking
    generalizations (PEP 448).
    """
    if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
        return False

    p = leaf.parent
    if p.type == syms.star_expr:
        # Star expressions are also used as assignment targets in extended
        # iterable unpacking (PEP 3132).  See what its parent is instead.
        if not p.parent:
            return False

        p = p.parent

    return p.type in within


def is_fstring(node: Node) -> bool:
    """Return True if the node is an f-string"""
    return node.type == syms.fstring


def fstring_tstring_to_string(node: Node) -> Leaf:
    """Converts an fstring or tstring node back to a string node."""
    string_without_prefix = str(node)[len(node.prefix) :]
    string_leaf = Leaf(token.STRING, string_without_prefix, prefix=node.prefix)
    string_leaf.lineno = node.get_lineno() or 0
    return string_leaf


def is_multiline_string(node: LN) -> bool:
    """Return True if `leaf` is a multiline string that actually spans many lines."""
    if isinstance(node, Node) and is_fstring(node):
        leaf = fstring_tstring_to_string(node)
    elif isinstance(node, Leaf):
        leaf = node
    else:
        return False

    return has_triple_quotes(leaf.value) and "\n" in leaf.value


def is_parent_function_or_class(node: Node) -> bool:
    assert node.type in {syms.suite, syms.simple_stmt}
    assert node.parent is not None
    # Note this works for suites / simple_stmts in async def as well
    return node.parent.type in {syms.funcdef, syms.classdef}


def is_stub_suite(node: Node) -> bool:
    """Return True if `node` is a suite with a stub body."""
    if node.parent is not None and not is_parent_function_or_class(node):
        return False

    # If there is a comment, we want to keep it.
    if node.prefix.strip():
        return False

    if (
        len(node.children) != 4
        or node.children[0].type != token.NEWLINE
        or node.children[1].type != token.INDENT
        or node.children[3].type != token.DEDENT
    ):
        return False

    if node.children[3].prefix.strip():
        return False

    return is_stub_body(node.children[2])


def is_stub_body(node: LN) -> bool:
    """Return True if `node` is a simple statement containing an ellipsis."""
    if not isinstance(node, Node) or node.type != syms.simple_stmt:
        return False

    if len(node.children) != 2:
        return False

    child = node.children[0]
    return (
        not child.prefix.strip()
        and child.type == syms.atom
        and len(child.children) == 3
        and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
    )


def is_atom_with_invisible_parens(node: LN) -> bool:
    """Given a `LN`, determines whether it's an atom `node` with invisible
    parens. Useful in dedupe-ing and normalizing parens.
    """
    if isinstance(node, Leaf) or node.type != syms.atom:
        return False

    first, last = node.children[0], node.children[-1]
    return (
        isinstance(first, Leaf)
        and first.type == token.LPAR
        and first.value == ""
        and isinstance(last, Leaf)
        and last.type == token.RPAR
        and last.value == ""
    )


def is_empty_par(leaf: Leaf) -> bool:
    return is_empty_lpar(leaf) or is_empty_rpar(leaf)


def is_empty_lpar(leaf: Leaf) -> bool:
    return leaf.type == token.LPAR and leaf.value == ""


def is_empty_rpar(leaf: Leaf) -> bool:
    return leaf.type == token.RPAR and leaf.value == ""


def is_import(leaf: Leaf) -> bool:
    """Return True if the given leaf starts an import statement."""
    p = leaf.parent
    t = leaf.type
    v = leaf.value
    return bool(
        t == token.NAME
        and (
            (v == "import" and p and p.type == syms.import_name)
            or (v == "from" and p and p.type == syms.import_from)
        )
    )


def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
    """Return True if the given leaf starts a with or async with statement."""
    return bool(
        leaf.type == token.NAME
        and leaf.value == "with"
        and leaf.parent
        and leaf.parent.type == syms.with_stmt
    ) or bool(
        leaf.type == token.ASYNC
        and leaf.next_sibling
        and leaf.next_sibling.type == syms.with_stmt
    )


def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
    """Return True if the given leaf starts an async def/for/with statement.

    Note that `async def` can be either an `async_stmt` or `async_funcdef`,
    the latter is used when it has decorators.
    """
    return bool(
        leaf.type == token.ASYNC
        and leaf.parent
        and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
    )


def is_type_comment(leaf: Leaf, mode: Mode) -> bool:
    """Return True if the given leaf is a type comment. This function should only
    be used for general type comments (excluding ignore annotations, which should
    use `is_type_ignore_comment`). Note that general type comments are no longer
    used in modern version of Python, this function may be deprecated in the future."""
    t = leaf.type
    v = leaf.value
    return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_comment_string(v, mode)


def is_type_comment_string(value: str, mode: Mode) -> bool:
    return value.startswith("#") and value[1:].lstrip().startswith("type:")


def is_type_ignore_comment(leaf: Leaf, mode: Mode) -> bool:
    """Return True if the given leaf is a type comment with ignore annotation."""
    t = leaf.type
    v = leaf.value
    return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(
        v, mode
    )


def is_type_ignore_comment_string(value: str, mode: Mode) -> bool:
    """Return True if the given string match with type comment with
    ignore annotation."""
    return is_type_comment_string(value, mode) and value.split(":", 1)[
        1
    ].lstrip().startswith("ignore")


def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
    """Wrap `child` in parentheses.

    This replaces `child` with an atom holding the parentheses and the old
    child.  That requires moving the prefix.

    If `visible` is False, the leaves will be valueless (and thus invisible).
    """
    lpar = Leaf(token.LPAR, "(" if visible else "")
    rpar = Leaf(token.RPAR, ")" if visible else "")
    prefix = child.prefix
    child.prefix = ""
    index = child.remove() or 0
    new_child = Node(syms.atom, [lpar, child, rpar])
    new_child.prefix = prefix
    parent.insert_child(index, new_child)


def unwrap_singleton_parenthesis(node: LN) -> LN | None:
    """Returns `wrapped` if `node` is of the shape ( wrapped ).

    Parenthesis can be optional. Returns None otherwise"""
    if len(node.children) != 3:
        return None

    lpar, wrapped, rpar = node.children
    if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
        return None

    return wrapped


def ensure_visible(leaf: Leaf) -> None:
    """Make sure parentheses are visible.

    They could be invisible as part of some statements (see
    :func:`normalize_invisible_parens` and :func:`visit_import_from`).
    """
    if leaf.type == token.LPAR:
        leaf.value = "("
    elif leaf.type == token.RPAR:
        leaf.value = ")"


def is_name_token(nl: NL) -> TypeGuard[Leaf]:
    return nl.type == token.NAME


def is_lpar_token(nl: NL) -> TypeGuard[Leaf]:
    return nl.type == token.LPAR


def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
    return nl.type == token.RPAR


def is_number_token(nl: NL) -> TypeGuard[Leaf]:
    return nl.type == token.NUMBER


def get_annotation_type(leaf: Leaf) -> Literal["return", "param", None]:
    """Returns the type of annotation this leaf is part of, if any."""
    ancestor = leaf.parent
    while ancestor is not None:
        if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
            return "return"
        if ancestor.parent and ancestor.parent.type == syms.tname:
            return "param"
        ancestor = ancestor.parent
    return None


def is_part_of_annotation(leaf: Leaf) -> bool:
    """Returns whether this leaf is part of a type annotation."""
    assert leaf.parent is not None
    return get_annotation_type(leaf) is not None


def first_leaf(node: LN) -> Leaf | None:
    """Returns the first leaf of the ancestor node."""
    if isinstance(node, Leaf):
        return node
    elif not node.children:
        return None
    else:
        return first_leaf(node.children[0])


def last_leaf(node: LN) -> Leaf | None:
    """Returns the last leaf of the ancestor node."""
    if isinstance(node, Leaf):
        return node
    elif not node.children:
        return None
    else:
        return last_leaf(node.children[-1])


def furthest_ancestor_with_last_leaf(leaf: Leaf) -> LN:
    """Returns the furthest ancestor that has this leaf node as the last leaf."""
    node: LN = leaf
    while node.parent and node.parent.children and node is node.parent.children[-1]:
        node = node.parent
    return node


def has_sibling_with_type(node: LN, type: int) -> bool:
    # Check previous siblings
    sibling = node.prev_sibling
    while sibling is not None:
        if sibling.type == type:
            return True
        sibling = sibling.prev_sibling

    # Check next siblings
    sibling = node.next_sibling
    while sibling is not None:
        if sibling.type == type:
            return True
        sibling = sibling.next_sibling

    return False
