/*
 * Copyright 2025 the original author or authors.
 * <p>
 * Licensed under the Moderne Source Available License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * https://docs.moderne.io/licensing/moderne-source-available-license
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import {JS} from "../tree";
import {JavaScriptVisitor} from "../visitor";
import {Comment, J, lastWhitespace, replaceLastWhitespace, Statement} from "../../java";
import {create as produce, Draft} from "mutative";
import {Cursor, isScope, Tree} from "../../tree";
import {BlankLinesStyle, getStyle, SpacesStyle, StyleKind, TabsAndIndentsStyle, WrappingAndBracesStyle} from "../style";
import {NamedStyles} from "../../style";
import {produceAsync} from "../../visitor";
import {Generator} from "../markers";
import {TabsAndIndentsVisitor} from "./tabs-and-indents-visitor";
import {NormalizeWhitespaceVisitor} from "./normalize-whitespace-visitor";
import {MinimumViableSpacingVisitor} from "./minimum-viable-spacing-visitor";
import {applyPrettierFormatting, getPrettierStyle} from "./prettier-format";

export {TabsAndIndentsVisitor} from "./tabs-and-indents-visitor";
export {NormalizeWhitespaceVisitor} from "./normalize-whitespace-visitor";
export {MinimumViableSpacingVisitor} from "./minimum-viable-spacing-visitor";

export const maybeAutoFormat = async <J2 extends J, P>(before: J2, after: J2, p: P, stopAfter?: J, parent?: Cursor): Promise<J2> => {
    if (before !== after) {
        return autoFormat(after, p, stopAfter, parent);
    }
    return after;
}

export const autoFormat = async <J2 extends J, P>(
    j: J2,
    p: P,
    stopAfter?: J,
    parent?: Cursor,
    styles?: NamedStyles<string>[]
): Promise<J2> =>
    (await new AutoformatVisitor(stopAfter, styles).visit(j, p, parent) as J2);

/**
 * Formats JavaScript/TypeScript code using a comprehensive set of formatting rules.
 *
 * Style resolution order (first match wins):
 * 1. Styles passed to the constructor
 * 2. Styles from source file markers (NamedStyles)
 * 3. IntelliJ defaults
 *
 * When a PrettierStyle is present (either in the styles array or as a marker on the source file),
 * Prettier is used for formatting. Otherwise, built-in formatting visitors are used.
 */
export class AutoformatVisitor<P> extends JavaScriptVisitor<P> {
    private readonly styles?: NamedStyles<string>[];

    constructor(private stopAfter?: Tree, styles?: NamedStyles<string>[]) {
        super();
        this.styles = styles;
    }

    async visit<R extends J>(tree: Tree, p: P, cursor?: Cursor): Promise<R | undefined> {
        // Check for PrettierStyle in styles array or as marker on source file
        // If found, delegate entirely to Prettier (skip other formatting visitors)
        const prettierStyle = getPrettierStyle(tree, cursor, this.styles);
        if (prettierStyle) {
            return applyPrettierFormatting(tree as R, prettierStyle, p, cursor, this.stopAfter);
        }

        const visitors = [
            new NormalizeWhitespaceVisitor(this.stopAfter),
            new MinimumViableSpacingVisitor(this.stopAfter),
            new BlankLinesVisitor(getStyle(StyleKind.BlankLinesStyle, tree, this.styles) as BlankLinesStyle, this.stopAfter),
            new WrappingAndBracesVisitor(getStyle(StyleKind.WrappingAndBracesStyle, tree, this.styles) as WrappingAndBracesStyle, this.stopAfter),
            new SpacesVisitor(getStyle(StyleKind.SpacesStyle, tree, this.styles) as SpacesStyle, this.stopAfter),
            new TabsAndIndentsVisitor(getStyle(StyleKind.TabsAndIndentsStyle, tree, this.styles) as TabsAndIndentsStyle, this.stopAfter),
        ]

        let t: R | undefined = tree as R;
        for (const visitor of visitors) {
            t = await visitor.visit(t, p, cursor);
            if (t === undefined) {
                return undefined;
            }
        }

        return t;
    }
}

export class SpacesVisitor<P> extends JavaScriptVisitor<P> {
    constructor(private style: SpacesStyle, private stopAfter?: Tree) {
        super();
    }

    override async visit<R extends J>(tree: Tree, p: P, parent?: Cursor): Promise<R | undefined> {
        if (this.cursor?.getNearestMessage("stop") != null) {
            return tree as R;
        }
        return super.visit(tree, p, parent);
    }

    override async postVisit(tree: J, p: P): Promise<J | undefined> {
        if (this.stopAfter != null && isScope(this.stopAfter, tree)) {
            this.cursor?.root.messages.set("stop", true);
        }
        return super.postVisit(tree, p);
    }

    protected async visitAlias(alias: JS.Alias, p: P): Promise<J | undefined> {
        const ret = await super.visitAlias(alias, p) as JS.Alias;
        return produce(ret, draft => {
            draft.propertyName.after.whitespace = " ";
        });
    }

    protected async visitArrayAccess(arrayAccess: J.ArrayAccess, p: P): Promise<J | undefined> {
        const ret = await super.visitArrayAccess(arrayAccess, p) as J.ArrayAccess;
        return produce(ret, draft => {
            // Preserve newlines - only modify if no newlines present
            if (!draft.dimension.index.element.prefix.whitespace.includes("\n")) {
                draft.dimension.index.element.prefix.whitespace = this.style.within.arrayBrackets ? " " : "";
            }
            if (!draft.dimension.index.after.whitespace.includes("\n")) {
                draft.dimension.index.after.whitespace = this.style.within.arrayBrackets ? " " : "";
            }
        });
    }

    protected async visitAssignment(assignment: J.Assignment, p: P): Promise<J | undefined> {
        const ret = await super.visitAssignment(assignment, p) as J.Assignment;
        return produce(ret, draft => {
            // Preserve newlines - only modify if no newlines present
            if (!draft.assignment.before.whitespace.includes("\n")) {
                draft.assignment.before.whitespace = this.style.aroundOperators.assignment ? " " : "";
            }
            if (!draft.assignment.element.prefix.whitespace.includes("\n")) {
                draft.assignment.element.prefix.whitespace = this.style.aroundOperators.assignment ? " " : "";
            }
        });
    }

    protected async visitBinary(binary: J.Binary, p: P): Promise<J | undefined> {
        const ret = await super.visitBinary(binary, p) as J.Binary;
        let property = false;
        switch (ret.operator.element.valueOf()) {
            case J.Binary.Type.And:
            case J.Binary.Type.Or:
                property = this.style.aroundOperators.logical
                break;
            case J.Binary.Type.Equal:
            case J.Binary.Type.NotEqual:
                property = this.style.aroundOperators.equality
                break;
            case J.Binary.Type.LessThan:
            case J.Binary.Type.LessThanOrEqual:
            case J.Binary.Type.GreaterThan:
            case J.Binary.Type.GreaterThanOrEqual:
                property = this.style.aroundOperators.relational
                break;
            case J.Binary.Type.BitAnd:
            case J.Binary.Type.BitOr:
            case J.Binary.Type.BitXor:
                property = this.style.aroundOperators.bitwise
                break;
            case J.Binary.Type.Addition:
            case J.Binary.Type.Subtraction:
                property = this.style.aroundOperators.additive
                break;
            case J.Binary.Type.Multiplication:
            case J.Binary.Type.Division:
            case J.Binary.Type.Modulo:
                property = this.style.aroundOperators.multiplicative
                break;
            case J.Binary.Type.LeftShift:
            case J.Binary.Type.RightShift:
            case J.Binary.Type.UnsignedRightShift:
                property = this.style.aroundOperators.shift
                break;
            default:
                throw new Error("Unsupported operator type " + ret.operator.element.valueOf());
        }
        return produce(ret, draft => {
            // Preserve newlines - only modify if no newlines present
            if (!draft.operator.before.whitespace.includes("\n")) {
                draft.operator.before.whitespace = property ? " " : "";
            }
            if (!draft.right.prefix.whitespace.includes("\n")) {
                draft.right.prefix.whitespace = property ? " " : "";
            }
        }) as J.Binary;
    }

    protected async visitCase(aCase: J.Case, p: P): Promise<J | undefined> {
        const ret = await super.visitCase(aCase, p) as J.Case;
        return ret && produce(ret, draft => {
            if (draft.caseLabels.elements[0].element.kind != J.Kind.Identifier || (draft.caseLabels.elements[0].element as J.Identifier).simpleName != "default") {
                // Preserve newlines - only set to space if no newline exists
                if (!draft.caseLabels.before.whitespace.includes("\n")) {
                    draft.caseLabels.before.whitespace = " ";
                }
            }
        });
    }

    protected async visitClassDeclaration(classDecl: J.ClassDeclaration, p: P): Promise<J | undefined> {
        const ret = await super.visitClassDeclaration(classDecl, p) as J.ClassDeclaration;
        // TODO typeParameters - IntelliJ doesn't seem to provide a setting for angleBrackets spacing for Typescript (while it does for Java),
        // thus we either introduce our own setting or just enforce the natural spacing with no setting

        return produce(ret, draft => {
            draft.body.prefix.whitespace = this.style.beforeLeftBrace.classInterfaceModuleLeftBrace ? " " : "";
        }) as J.ClassDeclaration;
    }

    public async visitContainer<T extends J>(container: J.Container<T>, p: P): Promise<J.Container<T>> {
        const ret = await super.visitContainer(container, p) as J.Container<T>;
        return produce(ret, draft => {
            if (draft.elements.length > 0) {
                // Apply beforeComma rule to all elements except the last
                // (last element's after is before closing bracket, not a comma)
                for (let i = 0; i < draft.elements.length - 1; i++) {
                    const afterWs = draft.elements[i].after.whitespace;
                    // Preserve newlines - only adjust when on same line
                    if (!afterWs.includes("\n")) {
                        draft.elements[i].after.whitespace = this.style.other.beforeComma ? " " : "";
                    }
                }
            }
            if (draft.elements.length > 1) {
                // Apply afterComma rule to elements after the first
                for (let i = 1; i < draft.elements.length; i++) {
                    const currentWs = draft.elements[i].element.prefix.whitespace;
                    // Preserve original newlines - only adjust spacing when elements are on same line
                    if (!currentWs.includes("\n")) {
                        draft.elements[i].element.prefix.whitespace = this.style.other.afterComma ? " " : "";
                    }
                }
            }
        });
    }

    protected async visitExportDeclaration(exportDeclaration: JS.ExportDeclaration, p: P): Promise<J | undefined> {
        const ret = await super.visitExportDeclaration(exportDeclaration, p) as JS.ExportDeclaration;
        return produce(ret, draft => {
            if (draft.exportClause) {
                draft.exportClause.prefix.whitespace = " ";
                if (draft.exportClause.kind == JS.Kind.NamedExports) {
                    const ne = (draft.exportClause as Draft<JS.NamedExports>);
                    if (ne.elements.elements.length > 0) {
                        // Check if this is a multi-line export (any element's prefix has a newline)
                        const isMultiLine = ne.elements.elements.some(e => e.element.prefix.whitespace.includes("\n"));
                        if (!isMultiLine) {
                            // Single-line: adjust brace spacing
                            ne.elements.elements[0].element.prefix.whitespace = this.style.within.es6ImportExportBraces ? " " : "";
                            ne.elements.elements[ne.elements.elements.length - 1].after.whitespace = this.style.within.es6ImportExportBraces ? " " : "";
                        } else {
                            // Multi-line: apply beforeComma rule to last element's after (for trailing commas)
                            // If it has only spaces (no newline), it's the space before a trailing comma
                            const lastAfter = ne.elements.elements[ne.elements.elements.length - 1].after.whitespace;
                            if (!lastAfter.includes("\n") && lastAfter.trim() === "") {
                                ne.elements.elements[ne.elements.elements.length - 1].after.whitespace = this.style.other.beforeComma ? " " : "";
                            }
                        }
                    }
                }
            }
            draft.typeOnly.before.whitespace = draft.typeOnly.element ? " " : "";
            if (draft.moduleSpecifier) {
                draft.moduleSpecifier.before.whitespace = " ";
                draft.moduleSpecifier.element.prefix.whitespace = " ";
            }
        })
    }

    protected async visitForLoop(forLoop: J.ForLoop, p: P): Promise<J | undefined> {
        const ret = await super.visitForLoop(forLoop, p) as J.ForLoop;
        return produceAsync(ret, async draft => {
            draft.control.prefix.whitespace = this.style.beforeParentheses.forParentheses ? " " : "";
            draft.control.init = await Promise.all(draft.control.init.map(async (oneInit, index) => {
                if (oneInit.element.kind === J.Kind.VariableDeclarations) {
                    const vd = oneInit.element as Draft<J.VariableDeclarations>;
                    if (vd.modifiers && vd.modifiers.length > 0) {
                        vd.modifiers[0].prefix.whitespace = "";
                    }
                }
                oneInit.after.whitespace = "";
                this.spaceBeforeRightPaddedElementDraft(oneInit, index === 0 ? this.style.within.forParentheses : true);
                return oneInit;
            }));
            if (draft.control.condition) {
                draft.control.condition.element.prefix.whitespace = " ";
                draft.control.condition.after.whitespace = this.style.other.beforeForSemicolon ? " " : "";
            }
            draft.control.update.forEach((oneUpdate, index) => {
                oneUpdate.element.prefix.whitespace = " ";
                oneUpdate.after.whitespace = (index === draft.control.update.length - 1 ? this.style.within.forParentheses : this.style.other.beforeForSemicolon) ? " " : "";
            });

            this.spaceBeforeRightPaddedElementDraft(draft.body, this.style.beforeLeftBrace.forLeftBrace);
            this.spaceAfterRightPaddedDraft(draft.body, false);
        });
    }

    protected async visitIf(iff: J.If, p: P): Promise<J | undefined> {
        const ret = await super.visitIf(iff, p) as J.If;
        return produceAsync(ret, async draft => {
            this.spaceBeforeDraft(draft.ifCondition, this.style.beforeParentheses.ifParentheses);
            this.spaceBeforeRightPaddedElementDraft(draft.ifCondition.tree, this.style.within.ifParentheses);
            this.spaceAfterRightPaddedDraft(draft.ifCondition.tree, this.style.within.ifParentheses);
            this.spaceBeforeRightPaddedElementDraft(draft.thenPart, this.style.beforeLeftBrace.ifLeftBrace);
            this.spaceAfterRightPaddedDraft(draft.thenPart, false);
        });
    }

    protected async visitImportDeclaration(jsImport: JS.Import, p: P): Promise<J | undefined> {
        const ret = await super.visitImportDeclaration(jsImport, p) as JS.Import;
        // Mutative detects same-value assignments and returns the original reference
        // when nothing actually changed, so no guard function is needed.
        return produce(ret, draft => {
            if (draft.importClause) {
                // Space after 'import' keyword always goes on importClause.prefix.
                // The parser is consistent here: both ImportClause and NamedBindings share
                // the same trivia span, but ImportClause consumes it first.
                draft.importClause.prefix.whitespace = " ";
                if (draft.importClause.name) {
                    // For import equals declarations (import X = Y), use assignment spacing
                    // For regular imports (import X from 'Y'), no space after name
                    draft.importClause.name.after.whitespace = draft.initializer
                        ? (this.style.aroundOperators.assignment ? " " : "")
                        : "";
                }
                if (draft.importClause.namedBindings) {
                    const hasDefaultImport = !!draft.importClause.name;
                    if (hasDefaultImport || draft.importClause.typeOnly) {
                        draft.importClause.namedBindings.prefix.whitespace = " ";
                    }
                    if (draft.importClause.namedBindings.kind == JS.Kind.NamedImports) {
                        const ni = draft.importClause.namedBindings as Draft<JS.NamedImports>;
                        const isMultiLine = ni.elements.elements.some(e => e.element.prefix.whitespace.includes("\n"));
                        if (!isMultiLine) {
                            const braceSpace = this.style.within.es6ImportExportBraces ? " " : "";
                            ni.elements.elements[0].element.prefix.whitespace = braceSpace;
                            ni.elements.elements[ni.elements.elements.length - 1].after.whitespace = braceSpace;
                        } else {
                            // Multi-line: apply beforeComma rule to last element's after (for trailing commas)
                            const lastAfter = ni.elements.elements[ni.elements.elements.length - 1].after;
                            if (!lastAfter.whitespace.includes("\n") && lastAfter.whitespace.trim() === "") {
                                lastAfter.whitespace = this.style.other.beforeComma ? " " : "";
                            }
                        }
                    }
                }
            }
            if (draft.moduleSpecifier) {
                draft.moduleSpecifier.before.whitespace = " ";
                draft.moduleSpecifier.element.prefix.whitespace = draft.importClause ? " " : "";
            }
        })
    }

    protected async visitIndexSignatureDeclaration(indexSignatureDeclaration: JS.IndexSignatureDeclaration, p: P): Promise<J | undefined> {
        const ret = await super.visitIndexSignatureDeclaration(indexSignatureDeclaration, p) as JS.IndexSignatureDeclaration;
        return produce(ret, draft => {
            draft.typeExpression.before.whitespace = this.style.other.beforeTypeReferenceColon ? " " : "";
            draft.typeExpression.element.prefix.whitespace = this.style.other.afterTypeReferenceColon ? " " : "";
        });
    }

    protected async visitMethodDeclaration(methodDecl: J.MethodDeclaration, p: P): Promise<J | undefined> {
        const ret = await super.visitMethodDeclaration(methodDecl, p) as J.MethodDeclaration;
        return produceAsync(ret, async draft => {
            if (draft.body) {
                this.spaceBeforeDraft(draft.body, this.style.beforeLeftBrace.functionLeftBrace);
            }

            if (draft.parameters.elements.length > 0 && draft.parameters.elements[0].element.kind != J.Kind.Empty) {
                draft.parameters.elements.forEach((param, index) => {
                    this.spaceAfterRightPaddedDraft(param, index === draft.parameters.elements.length - 1 ? this.style.within.functionDeclarationParentheses : this.style.other.beforeComma);
                    this.spaceBeforeDraft(param.element, index === 0 ? this.style.within.functionDeclarationParentheses : this.style.other.afterComma);
                    (param.element as Draft<J.VariableDeclarations>).variables[0].element.name.prefix.whitespace = "";
                    (param.element as Draft<J.VariableDeclarations>).variables[0].after.whitespace = "";
                });
            } else if (draft.parameters.elements.length == 1) {
                this.spaceBeforeRightPaddedElementDraft(draft.parameters.elements[0], this.style.within.functionDeclarationParentheses);
                this.spaceAfterRightPaddedDraft(draft.parameters.elements[0], false);
            }
            this.spaceBeforeContainerDraft(draft.parameters, this.style.beforeParentheses.functionDeclarationParentheses);

            // Handle generator asterisk spacing
            // - space before * is in the Generator marker's prefix
            // - space after * is in the method name's prefix
            const generatorIndex = ret.markers.markers.findIndex(m => m.kind === JS.Markers.Generator);
            if (generatorIndex >= 0) {
                const generator = draft.markers.markers[generatorIndex] as Draft<Generator>;
                generator.prefix.whitespace = this.style.other.beforeAsteriskInGenerator ? " " : "";
                draft.name.prefix.whitespace = this.style.other.afterAsteriskInGenerator ? " " : "";
            }

            // TODO typeParameters handling - see visitClassDeclaration
        });
    }

    protected async visitMethodInvocation(methodInv: J.MethodInvocation, p: P): Promise<J | undefined> {
        const ret = await super.visitMethodInvocation(methodInv, p) as J.MethodInvocation;
        return produceAsync(ret, async draft => {
            if (draft.select) {
                this.spaceBeforeContainerDraft(draft.arguments, this.style.beforeParentheses.functionCallParentheses);
            }
            if (ret.arguments.elements.length > 0 && ret.arguments.elements[0].element.kind != J.Kind.Empty) {
                draft.arguments.elements.forEach((arg, index) => {
                    this.spaceAfterRightPaddedDraft(arg, index === draft.arguments.elements.length - 1 ? this.style.within.functionCallParentheses : this.style.other.beforeComma);
                    this.spaceBeforeDraft(arg.element, index === 0 ? this.style.within.functionCallParentheses : this.style.other.afterComma);
                });
            } else if (ret.arguments.elements.length == 1) {
                this.spaceBeforeRightPaddedElementDraft(draft.arguments.elements[0], this.style.within.functionCallParentheses);
                this.spaceAfterRightPaddedDraft(draft.arguments.elements[0], false);
            }
            // TODO typeParameters handling - see visitClassDeclaration
        });
    }

    protected async visitPropertyAssignment(propertyAssignment: JS.PropertyAssignment, p: P): Promise<J | undefined> {
        const pa = await super.visitPropertyAssignment(propertyAssignment, p) as JS.PropertyAssignment;
        // Only adjust the space before the colon/equals if there's an initializer (not a shorthand property)
        if (pa.initializer) {
            return produceAsync(pa, draft => {
                draft.name.after.whitespace = this.style.other.beforePropertyNameValueSeparator ? " " : "";
            });
        }
        return pa;
    }

    protected async visitSwitch(switchNode: J.Switch, p: P): Promise<J | undefined> {
        const ret = await super.visitSwitch(switchNode, p) as J.Switch;
        return produceAsync(ret, async draft => {
            this.spaceBeforeDraft(draft.selector, this.style.beforeParentheses.switchParentheses);
            this.spaceBeforeRightPaddedElementDraft(draft.selector.tree, this.style.within.switchParentheses);
            this.spaceAfterRightPaddedDraft(draft.selector.tree, this.style.within.switchParentheses);
            this.spaceBeforeDraft(draft.cases, this.style.beforeLeftBrace.switchLeftBrace);

            for (const case_ of draft.cases.statements) {
                if (case_.element.kind === J.Kind.Case) {
                    (case_.element as Draft<J.Case>).caseLabels.elements[0].after.whitespace = "";
                }
            }
        });
    }

    protected async visitTernary(ternary: J.Ternary, p: P): Promise<J | undefined> {
        const ret = await super.visitTernary(ternary, p) as J.Ternary;
        return produceAsync(ret, async draft => {
            this.spaceBeforeLeftPaddedElementDraft(draft.truePart, this.style.ternaryOperator.beforeQuestionMark, this.style.ternaryOperator.afterQuestionMark);
            this.spaceBeforeLeftPaddedElementDraft(draft.falsePart, this.style.ternaryOperator.beforeColon, this.style.ternaryOperator.afterColon);
        });
    }

    protected async visitTry(try_: J.Try, p: P): Promise<J | undefined> {
        const ret = await super.visitTry(try_, p) as J.Try;
        return produceAsync(ret, async draft => {
            draft.body.prefix.whitespace = this.style.beforeLeftBrace.tryLeftBrace ? " " : "";
            draft.catches.forEach(catch_ => {
                this.spaceBeforeDraft(catch_, this.style.beforeKeywords.catchKeyword);
                catch_.parameter.prefix.whitespace = this.style.beforeParentheses.catchParentheses ? " " : "";
                this.spaceBeforeRightPaddedElementDraft(catch_.parameter.tree, this.style.within.catchParentheses);
                this.spaceAfterRightPaddedDraft(catch_.parameter.tree, this.style.within.catchParentheses);
                if (catch_.parameter.tree.element.variables.length > 0) {
                    catch_.parameter.tree.element.variables[catch_.parameter.tree.element.variables.length - 1].after.whitespace = "";
                }
                catch_.body.prefix.whitespace = this.style.beforeLeftBrace.catchLeftBrace ? " " : "";
            });
            if (draft.finally) {
                this.spaceBeforeLeftPaddedElementDraft(draft.finally, this.style.beforeKeywords.finallyKeyword, this.style.beforeLeftBrace.finallyLeftBrace);
            }
        });
    }

    protected async visitTypeDeclaration(typeDeclaration: JS.TypeDeclaration, p: P): Promise<J | undefined> {
        const ret = await super.visitTypeDeclaration(typeDeclaration, p) as JS.TypeDeclaration;
        return produce(ret, draft => {
            if (draft.modifiers.length > 0) {
                draft.name.before.whitespace = " ";
            }
            draft.name.element.prefix.whitespace = " ";
            // Preserve newlines - only modify if no newlines present
            if (!draft.initializer.before.whitespace.includes("\n")) {
                draft.initializer.before.whitespace = this.style.aroundOperators.assignment ? " " : "";
            }
            if (!draft.initializer.element.prefix.whitespace.includes("\n")) {
                draft.initializer.element.prefix.whitespace = this.style.aroundOperators.assignment ? " " : "";
            }
        });
    }

    protected async visitTypeInfo(typeInfo: JS.TypeInfo, p: P): Promise<J | undefined> {
        const ret = await super.visitTypeInfo(typeInfo, p) as JS.TypeInfo;
        return produceAsync(ret, async draft => {
            draft.prefix.whitespace = this.style.other.beforeTypeReferenceColon ? " " : "";
            this.spaceBeforeDraft(draft.typeIdentifier, this.style.other.afterTypeReferenceColon);
        });
    }

    protected async visitTypeLiteral(typeLiteral: JS.TypeLiteral, p: P): Promise<J | undefined> {
        const ret = await super.visitTypeLiteral(typeLiteral, p) as JS.TypeLiteral;
        // Apply objectLiteralTypeBraces spacing for single-line type literals
        if (ret.members && ret.members.statements.length > 0) {
            const stmts = ret.members.statements;
            const isSingleLine = !ret.members.end.whitespace.includes("\n") &&
                stmts.every(s => !s.element.prefix.whitespace.includes("\n"));
            if (isSingleLine) {
                return produce(ret, draft => {
                    const space = this.style.within.objectLiteralTypeBraces ? " " : "";
                    draft.members.statements[0].element.prefix.whitespace = space;
                    // For type literals, the space before } is in members.end, not in last statement's after
                    draft.members.end.whitespace = space;
                });
            }
        }
        return ret;
    }

    protected async visitUnary(unary: J.Unary, p: P): Promise<J | undefined> {
        const ret = await super.visitUnary(unary, p) as J.Unary;
        return produce(ret, draft => {
            const spacing = this.style.aroundOperators.unary;

            switch (draft.operator.element) {
                case J.Unary.Type.Not:
                    draft.expression.prefix.whitespace = this.style.aroundOperators.afterUnaryNotAndNotNull ? " " : "";
                    break;
                case J.Unary.Type.PreIncrement:
                case J.Unary.Type.PreDecrement:
                case J.Unary.Type.Negative:
                case J.Unary.Type.Positive:
                case J.Unary.Type.Complement:
                    draft.expression.prefix.whitespace = spacing ? " " : "";
                    break;

                case J.Unary.Type.PostIncrement:
                case J.Unary.Type.PostDecrement:
                    // postfix: don't add space after operand
                    break;
            }
        });
    }
    protected async visitVariable(variable: J.VariableDeclarations.NamedVariable, p: P): Promise<J | undefined> {
        const ret = await super.visitVariable(variable, p) as J.VariableDeclarations.NamedVariable;
        if (variable.initializer?.element?.kind == JS.Kind.StatementExpression
            && (variable.initializer.element as JS.StatementExpression).statement.kind == J.Kind.MethodDeclaration) {
            return ret;
        }
        return produceAsync(ret, async draft => {
            if (draft.initializer) {
                // Preserve newlines - only modify if no newlines present
                if (!draft.initializer.before.whitespace.includes("\n")) {
                    draft.initializer.before.whitespace = this.style.aroundOperators.assignment ? " " : "";
                }
                if (!draft.initializer.element.prefix.whitespace.includes("\n")) {
                    draft.initializer.element.prefix.whitespace = this.style.aroundOperators.assignment ? " " : "";
                }
            }
        });
    }

    protected async visitWhileLoop(whileLoop: J.WhileLoop, p: P): Promise<J | undefined> {
        const ret = await super.visitWhileLoop(whileLoop, p) as J.WhileLoop;
        return produceAsync(ret, async draft => {
            this.spaceBeforeRightPaddedElementDraft(draft.body, this.style.beforeLeftBrace.whileLeftBrace);
            this.spaceAfterRightPaddedDraft(draft.body, false);
            this.spaceBeforeDraft(draft.condition, this.style.beforeParentheses.whileParentheses);
            this.spaceBeforeRightPaddedElementDraft(draft.condition.tree, this.style.within.whileParentheses);
            this.spaceAfterRightPaddedDraft(draft.condition.tree, this.style.within.whileParentheses);
        });
    }

    protected async visitTypeParameter(typeParam: J.TypeParameter, p: P): Promise<J | undefined> {
        const ret = await super.visitTypeParameter(typeParam, p) as J.TypeParameter;
        return produce(ret, draft => {
            if (draft.bounds && draft.bounds.elements.length >= 2) {
                const constraintType = draft.bounds.elements[0];
                const defaultType = draft.bounds.elements[1];
                const hasConstraint = constraintType.element.kind !== J.Kind.Empty;
                const hasDefault = defaultType.element.kind !== J.Kind.Empty;

                if (hasConstraint) {
                    // Space before '=' for default type (in the `after` of constraint type)
                    if (hasDefault && !constraintType.after.whitespace.includes("\n")) {
                        constraintType.after.whitespace = this.style.aroundOperators.assignment ? " " : "";
                    }
                } else if (hasDefault) {
                    // No constraint, just default: space before '=' is in bounds.before
                    if (!draft.bounds.before.whitespace.includes("\n")) {
                        draft.bounds.before.whitespace = this.style.aroundOperators.assignment ? " " : "";
                    }
                }
            }
        });
    }

    /**
     * Modifies the after space of a RightPadded draft in place.
     */
    private spaceAfterRightPaddedDraft<T extends J>(draft: Draft<J.RightPadded<T>>, spaceAfter: boolean): void {
        if (draft.after.comments.length > 0) {
            // Perform the space rule for the suffix of the last comment only. Same as IntelliJ.
            SpacesVisitor.spaceLastCommentSuffixDraft(draft.after.comments, spaceAfter);
            return;
        }

        if (spaceAfter && SpacesVisitor.isNotSingleSpace(draft.after.whitespace)) {
            draft.after.whitespace = " ";
        } else if (!spaceAfter && SpacesVisitor.isOnlySpacesAndNotEmpty(draft.after.whitespace)) {
            draft.after.whitespace = "";
        }
    }

    /**
     * Modifies the before space of a LeftPadded draft and the element's prefix in place.
     */
    private spaceBeforeLeftPaddedElementDraft<T extends J>(draft: Draft<J.LeftPadded<T>>, spaceBeforePadding: boolean, spaceBeforeElement: boolean): void {
        if (draft.before.comments.length == 0) {
            // Preserve newlines - only modify if no newlines present
            if (!draft.before.whitespace.includes("\n")) {
                draft.before.whitespace = spaceBeforePadding ? " " : "";
            }
        }
        this.spaceBeforeDraft(draft.element, spaceBeforeElement);
    }

    /**
     * Modifies the element's prefix of a RightPadded draft in place.
     */
    private spaceBeforeRightPaddedElementDraft<T extends J>(draft: Draft<J.RightPadded<T>>, spaceBefore: boolean): void {
        this.spaceBeforeDraft(draft.element, spaceBefore);
    }

    /**
     * Modifies the prefix whitespace of a J draft in place.
     */
    private spaceBeforeDraft<T extends J>(draft: Draft<T>, spaceBefore: boolean): void {
        if (draft.prefix.comments.length > 0) {
            return;
        }
        if (spaceBefore && SpacesVisitor.isNotSingleSpace(draft.prefix.whitespace)) {
            draft.prefix.whitespace = " ";
        } else if (!spaceBefore && SpacesVisitor.isOnlySpacesAndNotEmpty(draft.prefix.whitespace)) {
            draft.prefix.whitespace = "";
        }
    }

    /**
     * Modifies the before space of a Container draft in place.
     */
    private spaceBeforeContainerDraft<T extends J>(draft: Draft<J.Container<T>>, spaceBefore: boolean): void {
        if (draft.before.comments.length > 0) {
            // Perform the space rule for the suffix of the last comment only. Same as IntelliJ.
            SpacesVisitor.spaceLastCommentSuffixDraft(draft.before.comments, spaceBefore);
            return;
        }

        if (spaceBefore && SpacesVisitor.isNotSingleSpace(draft.before.whitespace)) {
            draft.before.whitespace = " ";
        } else if (!spaceBefore && SpacesVisitor.isOnlySpacesAndNotEmpty(draft.before.whitespace)) {
            draft.before.whitespace = "";
        }
    }

    /**
     * Modifies the suffix of the last comment in an array of draft comments in place.
     */
    private static spaceLastCommentSuffixDraft(comments: Draft<Comment>[], spaceSuffix: boolean): void {
        this.spaceSuffixDraft(comments[comments.length - 1], spaceSuffix);
    }

    /**
     * Modifies the suffix of a Comment draft in place.
     */
    private static spaceSuffixDraft(draft: Draft<Comment>, spaceSuffix: boolean): void {
        if (spaceSuffix && this.isNotSingleSpace(draft.suffix)) {
            draft.suffix = " ";
        } else if (!spaceSuffix && this.isOnlySpacesAndNotEmpty(draft.suffix)) {
            draft.suffix = "";
        }
    }

    private static isOnlySpaces(str: string): boolean {
        return str.split('').every(char => char === ' ');
    }

    private static isOnlySpacesAndNotEmpty(s: string): boolean {
        return s !== "" && this.isOnlySpaces(s);
    }

    private static isNotSingleSpace(str: string): boolean {
        return this.isOnlySpaces(str) && str !== " ";
    }

    protected async visitNewClass(newClass: J.NewClass, p: P): Promise<J | undefined> {
        const ret = await super.visitNewClass(newClass, p) as J.NewClass;

        // Only handle object literals (NewClass with no class/constructor)
        if (ret.class) {
            return ret;
        }

        // Handle object literal brace spacing: { foo: 1 } vs {foo: 1}
        if (ret.body && ret.body.statements.length > 0) {
            return produce(ret, draft => {
                const stmts = draft.body!.statements;
                // Check if this is a multi-line object literal
                const isMultiLine = stmts.some(s => s.element.prefix.whitespace.includes("\n")) ||
                    draft.body!.end.whitespace.includes("\n");

                if (!isMultiLine) {
                    // Single-line: apply objectLiteralBraces spacing
                    const space = this.style.within.objectLiteralBraces ? " " : "";
                    stmts[0].element.prefix.whitespace = space;
                    draft.body!.end.whitespace = space;
                }
            });
        }
        return ret;
    }
}

export class WrappingAndBracesVisitor<P> extends JavaScriptVisitor<P> {
    constructor(private readonly style: WrappingAndBracesStyle, private stopAfter?: Tree) {
        super();
    }

    override async visit<R extends J>(tree: Tree, p: P, parent?: Cursor): Promise<R | undefined> {
        if (this.cursor?.getNearestMessage("stop") != null) {
            return tree as R;
        }
        return super.visit(tree, p, parent);
    }

    override async postVisit(tree: J, p: P): Promise<J | undefined> {
        if (this.stopAfter != null && isScope(this.stopAfter, tree)) {
            this.cursor?.root.messages.set("stop", true);
        }
        return super.postVisit(tree, p);
    }

    protected async visitVariableDeclarations(multiVariable: J.VariableDeclarations, p: P): Promise<J.VariableDeclarations> {
        const v = await super.visitVariableDeclarations(multiVariable, p) as J.VariableDeclarations;
        if (v.leadingAnnotations.length === 0) {
            return v;
        }
        const parent = this.cursor.parentTree()?.value;
        if (parent?.kind === J.Kind.Block) {
            return produce(v, draft => {
                draft.leadingAnnotations = this.withNewlines(draft.leadingAnnotations);
                if (draft.leadingAnnotations.length > 0) {
                    if (draft.modifiers.length > 0) {
                        draft.modifiers = this.withNewlineModifiers(draft.modifiers);
                    } else if (draft.typeExpression && !draft.typeExpression.prefix.whitespace.includes("\n")) {
                        draft.typeExpression.prefix.whitespace = "\n" + draft.typeExpression.prefix.whitespace;
                    }
                }
            });
        }
        return v;
    }

    protected async visitMethodDeclaration(method: J.MethodDeclaration, p: P): Promise<J.MethodDeclaration> {
        const m = await super.visitMethodDeclaration(method, p) as J.MethodDeclaration;
        if (m.leadingAnnotations.length === 0) {
            return m;
        }
        return produce(m, draft => {
            draft.leadingAnnotations = this.withNewlines(draft.leadingAnnotations);
            if (draft.leadingAnnotations.length > 0) {
                if (draft.modifiers.length > 0) {
                    draft.modifiers = this.withNewlineModifiers(draft.modifiers);
                } else if (draft.typeParameters && !draft.typeParameters.prefix.whitespace.includes("\n")) {
                    draft.typeParameters.prefix.whitespace = "\n" + draft.typeParameters.prefix.whitespace;
                } else if (draft.returnTypeExpression && !draft.returnTypeExpression.prefix.whitespace.includes("\n")) {
                    draft.returnTypeExpression.prefix.whitespace = "\n" + draft.returnTypeExpression.prefix.whitespace;
                } else if (!draft.name.prefix.whitespace.includes("\n")) {
                    draft.name.prefix.whitespace = "\n" + draft.name.prefix.whitespace;
                }
            }
        });
    }

    protected async visitElse(elsePart: J.If.Else, p: P): Promise<J.If.Else> {
        const e = await super.visitElse(elsePart, p) as J.If.Else;
        const hasBody = e.body.element.kind === J.Kind.Block || e.body.element.kind === J.Kind.If;

        if (!hasBody) {
            return e;
        }

        const shouldHaveNewline = this.style.ifStatement.elseOnNewLine;
        const hasNewline = e.prefix.whitespace.includes("\n");
        if ((shouldHaveNewline && !hasNewline) || (!shouldHaveNewline && hasNewline)) {
            return produce(e, draft => {
                if (shouldHaveNewline && !hasNewline) {
                    draft.prefix.whitespace = "\n" + draft.prefix.whitespace;
                } else {
                    draft.prefix.whitespace = "";
                }
            });
        }
        return e;
    }

    protected async visitClassDeclaration(classDecl: J.ClassDeclaration, p: P): Promise<J.ClassDeclaration> {
        const j = await super.visitClassDeclaration(classDecl, p) as J.ClassDeclaration;
        if (j.leadingAnnotations.length === 0) {
            return j;
        }
        return produce(j, draft => {
            draft.leadingAnnotations = this.withNewlines(draft.leadingAnnotations);
            if (draft.leadingAnnotations.length > 0) {
                if (draft.modifiers.length > 0) {
                    draft.modifiers = this.withNewlineModifiers(draft.modifiers);
                } else {
                    const kind = draft.classKind;
                    if (!kind.prefix.whitespace.includes("\n")) {
                        kind.prefix.whitespace = "\n" + kind.prefix.whitespace;
                    }
                }
            }
        });
    }

    protected async visitBlock(block: J.Block, p: P): Promise<J.Block> {
        const b = await super.visitBlock(block, p) as J.Block;
        const parentKind = this.cursor.parent?.value.kind;

        // Check if this is a "simple" block (empty or contains only a single J.Empty)
        const isSimpleBlock = b.statements.length === 0 ||
            (b.statements.length === 1 && b.statements[0].element.kind === J.Kind.Empty);

        // Object literals and type literals: always format empty ones as {} on single line
        if (parentKind === J.Kind.NewClass || parentKind === JS.Kind.TypeLiteral) {
            if (isSimpleBlock && this.blockIsMultiLine(b)) {
                return produce(b, draft => {
                    this.formatBlockOnOneLine(draft);
                });
            }
            return b;
        }

        if (isSimpleBlock) {
            const isMethodOrFunctionBody = parentKind === J.Kind.Lambda ||
                parentKind === J.Kind.MethodDeclaration;
            const keepInOneLine = isMethodOrFunctionBody
                ? this.style.keepWhenReformatting.simpleMethodsInOneLine
                : this.style.keepWhenReformatting.simpleBlocksInOneLine;

            if (keepInOneLine) {
                if (this.blockIsMultiLine(b)) {
                    return produce(b, draft => {
                        this.formatBlockOnOneLine(draft);
                    });
                }
            } else {
                if (!b.end.whitespace.includes("\n")) {
                    return produce(b, draft => {
                        draft.end = this.withNewlineSpace(draft.end);
                    });
                }
            }
        } else {
            // Non-simple blocks: ensure closing brace is on its own line
            if (!b.end.whitespace.includes("\n") && !b.statements[b.statements.length - 1].after.whitespace.includes("\n")) {
                return produce(b, draft => {
                    draft.end = this.withNewlineSpace(draft.end);
                });
            }
        }

        return b;
    }

    private blockIsMultiLine(b: J.Block): boolean {
        if (b.end.whitespace.includes("\n")) return true;
        if (b.statements.length === 1) {
            if (b.statements[0].element.prefix.whitespace.includes("\n")) return true;
            if (b.statements[0].after.whitespace.includes("\n")) return true;
        }
        return false;
    }

    private formatBlockOnOneLine(draft: Draft<J.Block>): void {
        if (draft.end.whitespace.includes("\n")) {
            draft.end.whitespace = draft.end.whitespace.replace(/\n\s*/g, "");
        }
        if (draft.statements.length === 1) {
            if (draft.statements[0].element.prefix.whitespace.includes("\n")) {
                draft.statements[0].element.prefix.whitespace = "";
            }
            if (draft.statements[0].after.whitespace.includes("\n")) {
                draft.statements[0].after.whitespace = "";
            }
        }
    }

    protected async visitSwitch(aSwitch: J.Switch, p: P): Promise<J | undefined> {
        return super.visitSwitch(aSwitch, p);
    }

    private withNewlines<T extends { prefix: J.Space }>(list: T[]): T[] {
        return list.map((a, index) => {
            if (index > 0 && !a.prefix.whitespace.includes("\n")) {
                return produce(a, draft => {
                    draft.prefix.whitespace = "\n" + draft.prefix.whitespace;
                });
            }
            return a;
        });
    }

    private withNewlineSpace(space: J.Space): J.Space {
        if (space.comments.length === 0) {
            space = produce(space, draft => {
               draft.whitespace = "\n" + space.whitespace;
            });
        }
        // TODO clarify the situation with Comment.isMultiline and then add this case
        // else if (space.comments[space.comments.length - 1].isMultiline) {
        //     space = space.withComments(ListUtils.mapLast(space.comments, c => c.withSuffix("\n")));
        // }

        return space;
    }

    private withNewlineModifiers(modifiers: J.Modifier[]): J.Modifier[] {
        if (modifiers.length === 0) {
            return modifiers;
        }
        const first = modifiers[0];
        if (!first.prefix.whitespace.includes("\n")) {
            return [
                produce(first, draft => {
                    draft.prefix.whitespace = "\n" + draft.prefix.whitespace;
                }),
                ...modifiers.slice(1),
            ];
        }
        return modifiers;
    }
}

export class BlankLinesVisitor<P> extends JavaScriptVisitor<P> {
    constructor(private readonly style: BlankLinesStyle, private stopAfter?: Tree) {
        super();
    }

    protected async preVisit(tree: J, p: P): Promise<J | undefined> {
        let ret = await super.preVisit(tree, p) as J;

        if (ret.kind === JS.Kind.CompilationUnit) {
            ret = produce(ret as JS.CompilationUnit, draft => {
                if (draft.prefix.comments.length == 0) {
                    draft.prefix.whitespace = "";
                }
            });
        }
        if (ret.kind === J.Kind.MethodDeclaration
            && this.cursor.parent?.parent?.parent?.value.kind === J.Kind.ClassDeclaration) {
            ret = produce(ret as JS.StatementExpression, draft => {
                this.ensurePrefixHasNewLine(draft);
            });
        }

        return ret;
    }

    override async visit<R extends J>(tree: Tree, p: P, cursor?: Cursor): Promise<R | undefined> {
        if (this.cursor?.getNearestMessage("stop") != null) {
            return tree as R;
        }
        return super.visit(tree, p, cursor);
    }

    protected async visitClassDeclaration(classDecl: J.ClassDeclaration, p: P): Promise<J.ClassDeclaration> {
        let ret = await super.visitClassDeclaration(classDecl, p) as J.ClassDeclaration;
        if (!ret.body) return ret;
        return produce(ret, draft => {
            const statements = draft.body.statements;
            if (statements.length > 0) {
                this.keepMaximumBlankLines(draft.body.statements[0].element, 0);

                const isInterface = draft.classKind.type === J.ClassDeclaration.Kind.Type.Interface;
                for (let i = 1; i < statements.length; i++) {
                    const previousElement = statements[i - 1].element;
                    let currentElement = statements[i].element;
                    if (previousElement.kind == J.Kind.VariableDeclarations || currentElement.kind == J.Kind.VariableDeclarations) {
                        const fieldBlankLines = isInterface
                            ? this.style.minimum.aroundFieldInInterface ?? 0
                            : this.style.minimum.aroundField;
                        this.minimumBlankLines(currentElement, fieldBlankLines);
                    }
                    if (previousElement.kind == J.Kind.MethodDeclaration || currentElement.kind == J.Kind.MethodDeclaration) {
                        const methodBlankLines = isInterface
                            ? this.style.minimum.aroundMethodInInterface ?? 0
                            : this.style.minimum.aroundMethod;
                        this.minimumBlankLines(currentElement, methodBlankLines);
                    }
                    this.keepMaximumBlankLines(currentElement, this.style.keepMaximum.inCode);
                    draft.body.statements[i].element = currentElement;
                }
            }

            const cu = this.cursor.firstEnclosing((x: any): x is JS.CompilationUnit => x.kind === JS.Kind.CompilationUnit) as JS.CompilationUnit | undefined;
            const topLevelIndex = cu?.statements.findIndex(s => s.element.id == draft.id);

            if (topLevelIndex !== undefined) {
                const isImportJustBeforeThis = topLevelIndex > 0 && cu?.statements[topLevelIndex - 1].element.kind === JS.Kind.Import;

                if (isImportJustBeforeThis) {
                    this.minimumBlankLines(draft, this.style.minimum.afterImports ?? 0);
                }
                if (topLevelIndex > 0) {
                    this.minimumBlankLines(draft, this.style.minimum.aroundClass);
                }
                this.keepMaximumBlankLines(draft, this.style.keepMaximum.inCode);
            }
        });
    }

    override async visitStatement(statement: Statement, p: P): Promise<Statement> {
        const ret = await super.visitStatement(statement, p) as Statement;
        const parent = this.cursor.parentTree()?.value;
        const grandparent = this.cursor.parentTree()?.parent?.value;

        return produce(ret, draft => {
            if (grandparent?.kind === J.Kind.ClassDeclaration && parent?.kind === J.Kind.Block) {
                const classDecl = grandparent as J.ClassDeclaration;
                const block = parent as J.Block;

                const isFirst = block.statements.length > 0 && block.statements[0].element.id === draft.id;
                if (!isFirst) {
                    if (draft.kind === J.Kind.VariableDeclarations) {
                        const declMax = classDecl.classKind.type === J.ClassDeclaration.Kind.Type.Interface
                            ? this.style.minimum.aroundFieldInInterface ?? 0
                            : this.style.minimum.aroundField;
                        this.minimumBlankLines(draft, declMax);
                    } else if (draft.kind === J.Kind.MethodDeclaration) {
                        const declMax = classDecl.classKind.type === J.ClassDeclaration.Kind.Type.Interface
                            ? this.style.minimum.aroundMethodInInterface ?? 0
                            : this.style.minimum.aroundMethod;
                        this.minimumBlankLines(draft, declMax);
                    } else if (draft.kind === J.Kind.Block) {
                        this.minimumBlankLines(draft, this.style.minimum.aroundFunction);
                    } else if (draft.kind === J.Kind.ClassDeclaration) {
                        this.minimumBlankLines(draft, this.style.minimum.aroundClass);
                    }
                    this.keepMaximumBlankLines(draft, this.style.keepMaximum.inCode);
                }
            } else if (parent?.kind === J.Kind.Block && grandparent?.kind !== J.Kind.NewClass && grandparent?.kind !== JS.Kind.TypeLiteral ||
                      (parent?.kind === JS.Kind.CompilationUnit && (parent! as JS.CompilationUnit).statements[0].element.id != draft.id) ||
                      (parent?.kind === J.Kind.Case && this.isInsideCaseStatements(parent as J.Case, draft))) {
                if (draft.kind != J.Kind.Case) {
                    this.ensurePrefixHasNewLine(draft);
                }
            }
        });
    }

    protected async visitBlock(block: J.Block, p: P): Promise<J.Block> {
        const b = await super.visitBlock(block, p) as J.Block;
        return produce(b, draft => {
            const parentKind = this.cursor.parent?.value.kind;
            // Skip newline for object literals (NewClass) and type literals (TypeLiteral)
            if (parentKind != J.Kind.NewClass && parentKind != JS.Kind.TypeLiteral) {
                draft.end = replaceLastWhitespace(draft.end, ws =>
                    ws.includes("\n") ? ws : ws.replace(/[ \t]+$/, '') + "\n"
                );
            }
        });
    }

    protected async visitEnumValue(enumValue: J.EnumValue, p: P): Promise<J.EnumValue> {
        const e = await super.visitEnumValue(enumValue, p) as J.EnumValue;
        this.keepMaximumBlankLines(e, this.style.keepMaximum.inCode);
        return e;
    }

    /**
     * Check if a statement is inside a Case's statements container or body,
     * not inside the caseLabels container.
     */
    private isInsideCaseStatements(caseNode: J.Case, statement: J): boolean {
        // Check if statement is in the statements container
        if (caseNode.statements.elements.some(s => s.element.id === statement.id)) {
            return true;
        }
        // Check if statement is the body
        if (caseNode.body?.element.id === statement.id) {
            return true;
        }
        return false;
    }

    private keepMaximumBlankLines<T extends J>(node: Draft<T>, max: number) {
        const whitespace = node.prefix.whitespace;
        const blankLines = BlankLinesVisitor.countNewlines(whitespace) - 1;
        if (blankLines > max) {
            let idx = 0;
            for (let i = 0; i < blankLines - max + 1; i++, idx++) {
                idx = whitespace.indexOf('\n', idx);
            }
            idx--;
            node.prefix.whitespace = whitespace.substring(idx);
        }
    }

    private minimumBlankLines<T extends J>(node: Draft<T>, min: number) {
        if (min <= 0) {
            return;
        }
        const currentNewlines = BlankLinesVisitor.countNewlines(node.prefix.whitespace);
        const needed = min - currentNewlines + 1;
        if (needed > 0) {
            node.prefix.whitespace = "\n".repeat(needed) + node.prefix.whitespace;
        }
    }

    private ensurePrefixHasNewLine(node: Draft<J>) {
        if (!node.prefix) return;

        // Check if newline already exists in the effective last whitespace
        if (lastWhitespace(node.prefix).includes("\n")) {
            return; // Already has a newline
        }

        if (node.kind === JS.Kind.ExpressionStatement) {
            this.ensurePrefixHasNewLine((node as JS.ExpressionStatement).expression);
        } else {
            node.prefix = replaceLastWhitespace(node.prefix, () => "\n");
        }
    }

    private static countNewlines(s: string): number {
        return [...s].filter(c => c === "\n").length;
    }

    override async postVisit(tree: J, p: P): Promise<J | undefined> {
        if (this.stopAfter != null && isScope(this.stopAfter, tree)) {
            this.cursor?.root.messages.set("stop", true);
        }
        return super.postVisit(tree, p);
    }
}

// Re-export prettier formatting utilities
export {prettierFormat} from "./prettier-format";