from __future__ import annotations

import sys
import textwrap
from enum import Enum, auto
from typing import TypeVar, Optional, cast, List

from rewrite.java import tree as j
from rewrite.java.support_types import J, Space, Comment, Statement
from rewrite.java.tree import (
    JRightPadded, JLeftPadded, JContainer, Block, ArrayDimension,
    ClassDeclaration, Empty, MethodDeclaration, MethodInvocation,
    FieldAccess, Identifier, Lambda, NewArray, Annotation, Literal,
    If, Binary, JavaSourceFile,
)
from rewrite.java.markers import TrailingComma
from rewrite.python import (
    PythonVisitor, TabsAndIndentsStyle, DictLiteral, CollectionLiteral,
    ExpressionStatement, OtherStyle, IntelliJ, ComprehensionExpression, PyComment,
    TrailingElseWrapper, Binary as PyBinary,
)
from rewrite.tree import Tree
from rewrite.visitor import P, T, Cursor
from rewrite.utils import list_map

J2 = TypeVar('J2', bound=J)


class TabsAndIndentsVisitor(PythonVisitor[P]):

    class IndentType(Enum):
        ALIGN = auto()
        INDENT = auto()
        CONTINUATION_INDENT = auto()

    def __init__(self, style: TabsAndIndentsStyle, other: OtherStyle = IntelliJ.other(),
                 stop_after: Optional[Tree] = None):
        self._stop_after = stop_after
        self._style = style
        self._other = other
        self._stop = False

    def visit(self, tree: Optional[Tree], p: P, parent: Optional[Cursor] = None) -> Optional[J]:
        if parent is not None:
            self._cursor = parent
            if tree is None:
                return cast(Optional[J], self.default_value(None, p))

            for c in parent.get_path_as_cursors():
                v = c.value
                space = None
                if isinstance(v, J):
                    space = v.prefix
                elif isinstance(v, JRightPadded):
                    space = v.after
                elif isinstance(v, JLeftPadded):
                    space = v.before
                elif isinstance(v, JContainer):
                    space = v.before

                if space is not None and '\n' in space.last_whitespace:
                    indent = self.find_indent(space)
                    if indent != 0:
                        c.put_message("last_indent", indent)

            for next_parent in parent.get_path():
                if isinstance(next_parent, J):
                    self.pre_visit(next_parent, p)
                    break

        return super().visit(tree, p)

    def pre_visit(self, tree: T, p: P) -> Optional[T]:
        if isinstance(tree, (JavaSourceFile, ArrayDimension, ClassDeclaration)):
            self.cursor.put_message("indent_type", self.IndentType.ALIGN)
        elif isinstance(tree, (Block, If)):
            self.cursor.put_message("indent_type", self.IndentType.INDENT)
        elif isinstance(tree, (DictLiteral, CollectionLiteral, NewArray, ComprehensionExpression)):
            self.cursor.put_message("indent_type", self.IndentType.CONTINUATION_INDENT
                if self._other.use_continuation_indent.collections_and_comprehensions else self.IndentType.INDENT)
        elif isinstance(tree, ExpressionStatement):
            pass
        elif isinstance(tree, (Binary, PyBinary)):
            if self.cursor.first_enclosing(j.Parentheses) is not None:
                self.cursor.put_message("indent_type", self.IndentType.ALIGN)
            else:
                self.cursor.put_message("indent_type", self.IndentType.INDENT)
        elif isinstance(tree, j.Expression):
            self.cursor.put_message("indent_type", self.IndentType.INDENT)

        return tree

    def post_visit(self, tree: T, p: P) -> Optional[T]:
        if self._stop_after and tree == self._stop_after:
            self._stop = True
        return tree

    # -------------------------------------------------------------------------
    # Core visit overrides (no loc parameter)
    # -------------------------------------------------------------------------

    def visit_space(self, space: Optional[Space], p: P) -> Space:
        if space is None:
            return space  # type: ignore

        space_context = self.cursor.get_nearest_message("space_context")
        parent = self.cursor.parent
        align_to_annotation = False

        if parent is not None:
            if isinstance(parent.value, Annotation):
                parent.parent_or_throw.put_message("after_annotation", True)
            elif not any(isinstance(_, Annotation) for _ in parent.get_path()):
                align_to_annotation = self.cursor.poll_nearest_message("after_annotation", False)

        # Detect file-start comments: comments in the prefix of a top-level element
        # with no preceding newline (parser places these in the statement prefix
        # rather than the CU prefix)
        if (space_context is None and space.comments and space.whitespace and
                "\n" not in space.whitespace and isinstance(self.cursor.value, J) and
                isinstance(self.cursor.parent_tree_cursor().value, JavaSourceFile)):
            space_context = "compilation_unit_prefix"

        indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0
        indent_type = self.cursor.parent_or_throw.get_nearest_message("indent_type") or self.IndentType.ALIGN

        if not space.comments and '\n' not in space.last_whitespace or parent is None:
            return space

        cursor_value = self.cursor.value

        align_block_prefix_to_parent = (
            space_context == "block_prefix" and '\n' in space.whitespace and
            isinstance(cursor_value, Block) and
            not isinstance(self.cursor.parent_tree_cursor().value, Block)
        )

        align_block_to_parent = space_context in ("catch_prefix", "try_finally", "else_prefix")

        if space_context == "extends" and "\n" in space.whitespace:
            indent_type = self.IndentType.CONTINUATION_INDENT

        if align_block_prefix_to_parent or align_block_to_parent or align_to_annotation:
            indent_type = self.IndentType.ALIGN

        if indent_type == self.IndentType.INDENT:
            indent += self._style.indent_size
        elif indent_type == self.IndentType.CONTINUATION_INDENT:
            indent += self._style.continuation_indent

        s = self._indent_to(space, indent, space_context)
        if isinstance(cursor_value, J):
            self.cursor.put_message("last_indent", indent)

        return s

    def visit_right_padded(self, right: Optional[JRightPadded], p: P) -> Optional[JRightPadded]:
        if right is None:
            return None

        self.cursor = Cursor(self._cursor, right)

        rp_context = self.cursor.parent_or_throw.get_message("rp_context", None)
        indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0

        t = right.element
        after = right.after

        if isinstance(t, J):
            elem = t
            trailing_comma = right.markers.find_first(TrailingComma)

            if '\n' in right.after.last_whitespace or '\n' in elem.prefix.last_whitespace:
                if rp_context == "method_declaration_parameter":
                    if isinstance(elem, Empty):
                        elem = elem.replace(prefix=self._indent_to(elem.prefix, indent))
                        after = right.after
                    else:
                        container = cast(JContainer, self.cursor.parent_or_throw.value)
                        elements = container.elements
                        first_arg = elements[0]
                        last_arg = elements[-1]

                        if self._style.method_declaration_parameters.align_multiline_parameters:
                            method = self.cursor.first_enclosing(MethodDeclaration)
                            if method is not None:
                                if "\n" in first_arg.prefix.last_whitespace:
                                    if self._other.use_continuation_indent.method_declaration_parameters:
                                        align_to = indent + self._style.continuation_indent
                                    else:
                                        align_to = self._get_length_of_whitespace(first_arg.prefix.last_whitespace)
                                else:
                                    align_to = indent + self.compute_first_parameter_offset(method, first_arg, self.cursor)
                                self.cursor.parent_or_throw.put_message("last_indent", align_to - self._style.continuation_indent)
                                elem = self.visit_and_cast(elem, J, p)
                                self.cursor.parent_or_throw.put_message("last_indent", indent)
                                after = self._indent_to(right.after, indent if t is last_arg else align_to)
                            else:
                                after = right.after
                        else:
                            elem = self.visit_and_cast(elem, J, p)
                            after = self._indent_to(right.after,
                                indent if t is last_arg else self._style.continuation_indent)

                elif rp_context == "method_invocation_argument":
                    elem, after = self._visit_method_invocation_argument_j_type(elem, right, indent, p)

                elif rp_context in ("collection_literal_element", "dict_literal_element"):
                    elem = self.visit_and_cast(elem, J, p)
                    args = cast(JContainer, self.cursor.parent_or_throw.value)
                    if not trailing_comma and args.padding.elements[-1] is right:
                        self.cursor.parent_or_throw.put_message("indent_type", self.IndentType.ALIGN)
                    after = self.visit_space(right.after, p)
                    if trailing_comma:
                        self.cursor.parent_or_throw.put_message("indent_type", self.IndentType.ALIGN)
                        trailing_comma = trailing_comma.replace(suffix=self.visit_space(trailing_comma.suffix, p))
                        right = right.replace(markers=right.markers.compute_by_type(TrailingComma, lambda _t: trailing_comma))

                elif rp_context is not None:
                    elem = self.visit_and_cast(elem, J, p)
                    after = self._indent_to(right.after, indent)

                else:
                    method_select_indent = self.cursor.get_nearest_message("method_select_indent")
                    elem = self.visit_and_cast(elem, J, p)
                    if method_select_indent is not None:
                        after = self._indent_to(right.after, method_select_indent)
                    else:
                        after = self.visit_space(right.after, p)

            else:
                if rp_context in ("method_invocation_argument",):
                    any_other_arg_on_own_line = False
                    if "\n" not in elem.prefix.last_whitespace:
                        args = cast(JContainer, self.cursor.parent_or_throw.value)
                        for arg in args.padding.elements:
                            if arg == self.cursor.value:
                                continue
                            if "\n" in arg.element.prefix.last_whitespace:
                                any_other_arg_on_own_line = True
                                break
                        if not any_other_arg_on_own_line:
                            elem = self.visit_and_cast(elem, J, p)
                            after = self._indent_to(right.after, indent)

                    if not any_other_arg_on_own_line:
                        if not isinstance(elem, Binary):
                            if not isinstance(elem, MethodInvocation) or "\n" in elem.prefix.last_whitespace:
                                self.cursor.put_message("last_indent", indent + self._style.continuation_indent)
                            else:
                                select = elem.select
                                if isinstance(select, (FieldAccess, Identifier, MethodInvocation)):
                                    self.cursor.put_message("last_indent", indent + self._style.continuation_indent)

                        elem = self.visit_and_cast(elem, J, p)
                        after = self.visit_space(right.after, p)
                else:
                    elem = self.visit_and_cast(elem, J, p)
                    after = self.visit_space(right.after, p)

            t = cast(T, elem)
        else:
            after = self.visit_space(right.after, p)

        self.cursor = self.cursor.parent  # type: ignore
        if t is right.element and after is right.after:
            return right
        return right.replace(element=t, after=after)

    def visit_container(self, container: Optional[JContainer], p: P) -> Optional[JContainer]:
        if container is None:
            return container

        self._cursor = Cursor(self._cursor, container)

        parent_tree = self.cursor.parent_tree_cursor().value

        rp_context = None
        needs_continuation_on_before = False
        if isinstance(parent_tree, MethodDeclaration):
            rp_context = "method_declaration_parameter"
            needs_continuation_on_before = True
        elif isinstance(parent_tree, (MethodInvocation, Annotation)):
            rp_context = "method_invocation_argument"
        elif isinstance(parent_tree, DictLiteral):
            rp_context = "dict_literal_element"
        elif isinstance(parent_tree, CollectionLiteral):
            rp_context = "collection_literal_element"
        elif isinstance(parent_tree, ClassDeclaration):
            needs_continuation_on_before = True

        self.cursor.put_message("rp_context", rp_context)

        indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0

        if '\n' in container.before.last_whitespace:
            if needs_continuation_on_before:
                before = self._indent_to(container.before, indent + self._style.continuation_indent)
                self.cursor.put_message("indent_type", self.IndentType.ALIGN)
                self.cursor.put_message("last_indent", indent + self._style.continuation_indent)
            else:
                before = self.visit_space(container.before, p)
            js = list_map(lambda t: self.visit_right_padded(t, p), container.padding.elements)
        else:
            if isinstance(parent_tree, MethodDeclaration):
                self.cursor.put_message("indent_type",
                    self.IndentType.CONTINUATION_INDENT if self._other.use_continuation_indent.method_declaration_parameters
                    else self.IndentType.INDENT)
            elif isinstance(parent_tree, (MethodInvocation, Annotation)):
                self.cursor.put_message("indent_type",
                    self.IndentType.CONTINUATION_INDENT if self._other.use_continuation_indent.method_call_arguments
                    else self.IndentType.INDENT)
            elif isinstance(parent_tree, ClassDeclaration):
                self.cursor.put_message("indent_type", self.IndentType.CONTINUATION_INDENT)
            before = self.visit_space(container.before, p)
            js = list_map(lambda t: self.visit_right_padded(t, p), container.padding.elements)

        self._cursor = self._cursor.parent  # type: ignore

        if container.padding.elements is js and container.before is before:
            return container
        return container.replace(before=before).padding.replace(elements=js)

    # -------------------------------------------------------------------------
    # Visit method overrides for setting space_context
    # -------------------------------------------------------------------------

    def visit_block(self, block: Block, p: P) -> J:
        self.cursor.put_message("space_context", "block_prefix")
        block = block.replace(prefix=self.visit_space(block.prefix, p))
        self.cursor.put_message("space_context", None)

        temp_stmt = self.visit_statement(block, p)
        if not isinstance(temp_stmt, Block):
            return temp_stmt
        block = temp_stmt

        block = block.replace(markers=self.visit_markers(block.markers, p))
        block = block.padding.replace(static=self.visit_right_padded(block.padding.static, p))
        block = block.padding.replace(statements=list_map(
            lambda stmt: self.visit_right_padded(stmt, p), block.padding.statements))

        self.cursor.put_message("space_context", "block_end")
        block = block.replace(end=self.visit_space(block.end, p))
        self.cursor.put_message("space_context", None)

        return block

    def visit_catch(self, catch: j.Try.Catch, p: P) -> J:
        self.cursor.put_message("space_context", "catch_prefix")
        catch = catch.replace(prefix=self.visit_space(catch.prefix, p))
        self.cursor.put_message("space_context", None)

        catch = catch.replace(markers=self.visit_markers(catch.markers, p))
        catch = catch.replace(
            parameter=self.visit_and_cast(catch.parameter, j.ControlParentheses, p))
        catch = catch.replace(body=cast(Block, self.visit(catch.body, p)))
        return catch

    def visit_try(self, try_: j.Try, p: P) -> J:
        try_ = try_.replace(prefix=self.visit_space(try_.prefix, p))

        temp_stmt = self.visit_statement(try_, p)
        if not isinstance(temp_stmt, j.Try):
            return temp_stmt
        try_ = temp_stmt

        try_ = try_.replace(markers=self.visit_markers(try_.markers, p))
        try_ = try_.replace(resources=self.visit_container(try_.padding.resources, p))
        try_ = try_.replace(body=self.visit_and_cast(try_.body, Block, p))
        try_ = try_.replace(catches=list_map(
            lambda c: self.visit_and_cast(c, j.Try.Catch, p), try_.catches))

        if try_.padding.finally_:
            self.cursor.put_message("space_context", "try_finally")
            try_ = try_.replace(finally_=self.visit_left_padded(try_.padding.finally_, p))
            self.cursor.put_message("space_context", None)

        return try_

    def visit_trailing_else_wrapper(self, wrapper: TrailingElseWrapper, p: P) -> J:
        wrapper = wrapper.replace(prefix=self.visit_space(wrapper.prefix, p))

        temp_stmt = cast(Statement, self.visit_statement(wrapper, p))
        if not isinstance(temp_stmt, type(wrapper)):
            return temp_stmt
        wrapper = temp_stmt

        wrapper = wrapper.replace(markers=self.visit_markers(wrapper.markers, p))
        wrapper = wrapper.replace(
            statement=self.visit_and_cast(wrapper.statement, Statement, p)
        )

        self.cursor.put_message("space_context", "else_prefix")
        wrapper = wrapper.padding.replace(
            else_block=self.visit_left_padded(wrapper.padding.else_block, p)
        )
        self.cursor.put_message("space_context", None)
        return wrapper

    def visit_else(self, else_: j.If.Else, p: P) -> J:
        self.cursor.put_message("space_context", "else_prefix")
        else_ = else_.replace(prefix=self.visit_space(else_.prefix, p))
        self.cursor.put_message("space_context", None)

        else_ = else_.replace(markers=self.visit_markers(else_.markers, p))
        else_ = else_.padding.replace(body=self.visit_right_padded(else_.padding.body, p))
        return else_

    def visit_compilation_unit(self, cu, p: P) -> J:
        self.cursor.put_message("space_context", "compilation_unit_prefix")
        cu = cu.replace(prefix=self.visit_space(cu.prefix, p))
        self.cursor.put_message("space_context", None)

        cu = cu.replace(markers=self.visit_markers(cu.markers, p))
        cu = cu.padding.replace(imports=list_map(
            lambda v: self.visit_right_padded(v, p), cu.padding.imports))
        cu = cu.padding.replace(statements=list_map(
            lambda v: self.visit_right_padded(v, p), cu.padding.statements))
        cu = cu.replace(eof=self.visit_space(cu.eof, p))
        return cu

    def visit_class_declaration(self, class_decl: ClassDeclaration, p: P) -> J:
        self.cursor.put_message("space_context", "class_declaration_prefix")
        class_decl = class_decl.replace(prefix=self.visit_space(class_decl.prefix, p))
        self.cursor.put_message("space_context", None)

        temp_stmt = self.visit_statement(class_decl, p)
        if not isinstance(temp_stmt, ClassDeclaration):
            return temp_stmt
        class_decl = temp_stmt

        class_decl = class_decl.replace(markers=self.visit_markers(class_decl.markers, p))
        class_decl = class_decl.replace(leading_annotations=list_map(
            lambda a: self.visit_and_cast(a, Annotation, p), class_decl.leading_annotations))
        class_decl = class_decl.replace(modifiers=list_map(
            lambda m: self.visit_and_cast(m, j.Modifier, p), class_decl.modifiers))
        class_decl = class_decl.replace(kind=self.visit_and_cast(
            class_decl.padding.kind, j.ClassDeclaration.Kind, p))
        class_decl = class_decl.replace(name=self.visit_and_cast(
            class_decl.name, j.Identifier, p))
        class_decl = class_decl.replace(type_parameters=self.visit_container(
            class_decl.padding.type_parameters, p))
        class_decl = class_decl.replace(primary_constructor=self.visit_container(
            class_decl.padding.primary_constructor, p))

        self.cursor.put_message("space_context", "extends")
        class_decl = class_decl.padding.replace(
            extends=self.visit_left_padded(class_decl.padding.extends, p))
        self.cursor.put_message("space_context", None)

        class_decl = class_decl.replace(implements=self.visit_container(
            class_decl.padding.implements, p))
        class_decl = class_decl.replace(permits=self.visit_container(
            class_decl.padding.permits, p))
        class_decl = class_decl.replace(body=self.visit_and_cast(
            class_decl.body, Block, p))
        return class_decl

    def visit_parentheses(self, parens: j.Parentheses, p: P) -> J:
        parens = parens.replace(prefix=self.visit_space(parens.prefix, p))
        temp_expr = self.visit_expression(parens, p)
        if not isinstance(temp_expr, j.Parentheses):
            return temp_expr
        parens = temp_expr
        parens = parens.replace(markers=self.visit_markers(parens.markers, p))

        tree_rp = parens.padding.tree
        elem = tree_rp.element

        # When a Binary expression starts on the same line as '(', compute
        # its column and set last_indent so continuation lines align there
        if isinstance(elem, (Binary, PyBinary)) and '\n' not in elem.prefix.last_whitespace:
            col = self._compute_column_of(elem)
            if col >= 0:
                self.cursor.put_message("last_indent", col)

        if isinstance(elem, J):
            elem = self.visit_and_cast(elem, J, p)

        after = tree_rp.after
        if '\n' in after.last_whitespace:
            # Close paren should align to the enclosing indent level
            indent = cast(int, self.cursor.parent_or_throw.get_nearest_message("last_indent")) or 0
            after = self._indent_to(after, indent)
        else:
            after = self.visit_space(after, p)

        if elem is not tree_rp.element or after is not tree_rp.after:
            parens = parens.padding.replace(tree=tree_rp.replace(element=elem, after=after))
        return parens

    def visit_binary(self, binary: Binary, p: P) -> J:
        if self.cursor.first_enclosing(j.Parentheses) is not None:
            # Inside parentheses: operator before spaces align to same indent as first line
            binary = binary.replace(prefix=self.visit_space(binary.prefix, p))
            temp_expr = self.visit_expression(binary, p)
            if not isinstance(temp_expr, Binary):
                return temp_expr
            binary = temp_expr
            binary = binary.replace(markers=self.visit_markers(binary.markers, p))
            binary = binary.replace(left=self.visit_and_cast(binary.left, j.Expression, p))

            op = binary.padding.operator
            if '\n' in op.before.last_whitespace:
                indent = cast(int, self.cursor.get_nearest_message("last_indent")) or 0
                before = self._indent_to(op.before, indent)
                if before is not op.before:
                    binary = binary.padding.replace(operator=op.replace(before=before))
            else:
                binary = binary.padding.replace(
                    operator=self.visit_left_padded(binary.padding.operator, p))

            binary = binary.replace(right=self.visit_and_cast(binary.right, j.Expression, p))
            return binary
        return super().visit_binary(binary, p)

    def visit_method_invocation(self, method: MethodInvocation, p: P) -> J:
        select = method.padding.select
        if select is not None and '\n' in select.after.last_whitespace:
            # Skip column alignment for method chains inside parentheses.
            # A chain exists when the method's select is itself a
            # MethodInvocation, or when there's an enclosing
            # MethodInvocation between this node and the Parentheses.
            skip = isinstance(select.element, MethodInvocation)
            if not skip:
                for c in self.cursor.get_path_as_cursors():
                    v = c.value
                    if v is method:
                        continue
                    if isinstance(v, j.Parentheses):
                        break
                    if isinstance(v, MethodInvocation):
                        skip = True
                        break
            if not skip:
                col = self._compute_select_column(method)
                if col >= 0:
                    self.cursor.put_message("method_select_indent", col)
        return super().visit_method_invocation(method, p)

    def _compute_column_of(self, target: Tree) -> int:
        """Compute the column position of a target tree element by printing from line start."""
        from rewrite.python.printer import PythonPrinter, PrintOutputCapture

        line_start = None
        line_start_cursor = None
        for c in self.cursor.get_path_as_cursors():
            v = c.value
            if isinstance(v, J):
                line_start = v
                line_start_cursor = c
                if '\n' in v.prefix.whitespace:
                    break
        if line_start is None:
            return -1

        class _ColumnCounter(PrintOutputCapture):
            def __init__(self):
                super().__init__()
                self.col = 0
                self.found = False

            def append(self, text):
                if text and not self.found:
                    for ch in text:
                        self.col = 0 if ch == '\n' else self.col + 1
                return self

        class _Printer(PythonPrinter):
            def __init__(self):
                super().__init__()
                orig_visit = self._delegate.visit
                def _check(tree, p):
                    if tree is target:
                        p.found = True
                        return tree
                    return orig_visit(tree, p) if not p.found else tree
                self._delegate.visit = _check

            def visit(self, tree, p, parent=None):
                if p.found or tree is target:
                    p.found = True
                    return tree
                return super().visit(tree, p)

        counter = _ColumnCounter()
        printer = _Printer()
        # Build initial cursor with proper J parent so context-dependent
        # printing (e.g., = vs :=) works correctly
        initial = Cursor(None, Cursor.ROOT_VALUE)
        assert line_start_cursor is not None
        for c in line_start_cursor.get_path_as_cursors():
            v = c.value
            if v is line_start:
                continue
            if isinstance(v, J):
                initial = Cursor(Cursor(None, Cursor.ROOT_VALUE), v)
                break
        printer.set_cursor(initial)
        printer.visit(line_start, counter)
        return counter.col if counter.found else -1

    def _compute_select_column(self, method: MethodInvocation) -> int:
        from rewrite.python.printer import PythonPrinter, PrintOutputCapture

        line_start = None
        for c in self.cursor.get_path_as_cursors():
            v = c.value
            if isinstance(v, J):
                line_start = v
                if '\n' in v.prefix.whitespace:
                    break
        if line_start is None:
            return -1

        select = method.select

        class _ColumnCounter(PrintOutputCapture):
            def __init__(self):
                super().__init__()
                self.col = 0
                self.found = False

            def append(self, text):
                if text and not self.found:
                    for ch in text:
                        self.col = 0 if ch == '\n' else self.col + 1
                return self

        class _Printer(PythonPrinter):
            def __init__(self, target):
                super().__init__()
                orig_visit = self._delegate.visit
                def _check(tree, p):
                    if tree is target:
                        p.found = True
                        return tree
                    return orig_visit(tree, p) if not p.found else tree
                self._delegate.visit = _check

            def visit(self, tree, p, parent=None):
                if p.found or tree is select:
                    p.found = True
                    return tree
                return super().visit(tree, p)

        counter = _ColumnCounter()
        _Printer(select).print(line_start, counter)
        return counter.col if counter.found else -1

    # -------------------------------------------------------------------------
    # Expression statement (docstring alignment)
    # -------------------------------------------------------------------------

    @staticmethod
    def _is_doc_comment(expression_statement: ExpressionStatement, cursor: Cursor) -> bool:
        expr = expression_statement.expression
        return isinstance(expr, Literal) and isinstance(expr.value_source, str) and (
            (expr.value_source.startswith('"""') and expr.value_source.endswith('"""')) or
            (expr.value_source.startswith("'''") and expr.value_source.endswith("'''"))) and \
            cursor.first_enclosing(Block) is not None

    def visit_expression_statement(self, expr_stmt: ExpressionStatement, p: P) -> J:
        if self._is_doc_comment(expr_stmt, self.cursor):
            prefix_before = len(expr_stmt.prefix.last_whitespace.split("\n")[-1])
            stm = cast(ExpressionStatement, super().visit_expression_statement(expr_stmt, p))
            literal = cast(Literal, stm.expression)
            shift = len(stm.prefix.last_whitespace.split("\n")[-1]) - prefix_before
            return stm.replace(expression=
                literal.replace(value_source=textwrap.indent(str(literal.value_source), shift * " ")[shift:]))

        return super().visit_expression_statement(expr_stmt, p)

    # -------------------------------------------------------------------------
    # First parameter offset calculation
    # -------------------------------------------------------------------------

    @staticmethod
    def compute_first_parameter_offset(method: MethodDeclaration, first_arg: J, cursor: Cursor) -> int:
        from rewrite.python.printer import PythonPrinter
        method = method.replace(leading_annotations=[])
        source = PythonPrinter().print(method)
        paren_idx = source.index("(")
        def_idx = source.index("def")
        async_idx = source.find("async")
        start_idx = async_idx if async_idx != -1 and async_idx < def_idx else def_idx
        return paren_idx - start_idx + 1 + len(first_arg.prefix.last_whitespace)

    # -------------------------------------------------------------------------
    # Helpers
    # -------------------------------------------------------------------------

    def _visit_method_invocation_argument_j_type(self, elem: J, right: JRightPadded,
                                                  indent: int, p: P) -> tuple:
        if "\n" not in elem.prefix.last_whitespace and isinstance(elem, Lambda):
            body = elem.body
            if not isinstance(body, Binary):
                if "\n" not in body.prefix.last_whitespace:
                    self.cursor.parent_or_throw.put_message("last_indent", indent + self._style.continuation_indent)

        elem = self.visit_and_cast(elem, J, p)
        assert elem is not None
        after = self._indent_to(right.after, indent)
        if after.comments or "\n" in after.last_whitespace:
            parent = self.cursor.parent_tree_cursor()
            grandparent = parent.parent_tree_cursor()
            if isinstance(grandparent.value, MethodInvocation) and grandparent.value.select == parent.value:
                grandparent.put_message("last_indent", indent)
                grandparent.put_message("chained_indent", indent)
        return elem, after

    def _indent_to(self, space: Space, column: int, context: Optional[str] = None) -> Space:
        s = space
        whitespace = s.whitespace

        if context == "compilation_unit_prefix" and whitespace:
            s = s.replace(whitespace="")
        elif not s.comments and "\n" not in s.last_whitespace:
            return s

        if not s.comments:
            indent = self.find_indent(s)
            if indent != column:
                shift = column - indent
                s = s.replace(whitespace=self._indent(whitespace, shift))
        else:
            def whitespace_indent(text: str) -> str:
                result: List[str] = []
                for c in text:
                    if c == '\n' or c == '\r':
                        return ''.join(result)
                    elif c.isspace():
                        result.append(c)
                    else:
                        return ''.join(result)
                return ''.join(result)

            has_file_leading_comment = space.comments and (
                context == "compilation_unit_prefix" or
                context == "block_end" or
                (context == "class_declaration_prefix" and space.comments[0].multiline)
            )

            final_column = column + self._style.indent_size if context == "block_end" else column
            last_indent_str: str = space.whitespace[space.whitespace.rfind('\n') + 1:]
            indent = self._get_length_of_whitespace(whitespace_indent(last_indent_str))

            if indent != final_column or s.comments:
                if (has_file_leading_comment or ("\n" in whitespace)) and (
                        not (s.comments and isinstance(s.comments[0], PyComment) and
                             not s.comments[0].multiline and self._get_length_of_whitespace(space.whitespace) == 0)):
                    shift = final_column - indent
                    s = s.replace(whitespace=whitespace[:whitespace.rfind('\n') + 1] + self._indent(last_indent_str, shift))

                final_space = s
                last_comment_pos = len(s.comments) - 1

                def _process_comment(idx: int, c: Comment) -> Comment:
                    if isinstance(c, PyComment) and not c.multiline:
                        if idx != last_comment_pos and self._get_length_of_whitespace(c.suffix) == 0:
                            return c

                    prior_suffix = space.whitespace if idx == 0 else final_space.comments[idx - 1].suffix

                    if context == "block_end" and idx != len(final_space.comments) - 1:
                        to_column = column + self._style.indent_size
                    else:
                        to_column = column

                    new_c = c
                    if "\n" in prior_suffix or has_file_leading_comment:
                        new_c = c

                    if '\n' in new_c.suffix:
                        suffix_indent = self._get_length_of_whitespace(new_c.suffix)
                        shift = to_column - suffix_indent
                        new_c = new_c.replace(suffix=self._indent(new_c.suffix, shift))

                    return new_c

                s = s.replace(comments=list_map(lambda c, i: _process_comment(i, c), s.comments))
        return s

    def _indent(self, whitespace: str, shift: int) -> str:
        return self._shift(whitespace, shift)

    def _shift(self, text: str, shift: int) -> str:
        tab_indent = self._style.tab_size
        if not self._style.use_tab_character:
            tab_indent = sys.maxsize

        if shift > 0:
            text += '\t' * (shift // tab_indent)
            text += ' ' * (shift % tab_indent)
        else:
            if self._style.use_tab_character:
                len_text = len(text) + (shift // tab_indent)
            else:
                len_text = len(text) + shift
            if len_text >= 0:
                text = text[:len_text]

        return text

    def find_indent(self, space: Space) -> int:
        return self._get_length_of_whitespace(space.indent)

    def _get_length_of_whitespace(self, whitespace: Optional[str]) -> int:
        if whitespace is None:
            return 0
        length = 0
        for c in whitespace:
            length += self._style.tab_size if c == '\t' else 1
            if c in ('\n', '\r'):
                length = 0
        return length
