diff --git a/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java b/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java
index 7b625828e83d..b22643def8e3 100644
--- a/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java
+++ b/java/java-impl/src/com/intellij/codeInspection/EnhancedSwitchMigrationInspection.java
@@ -512,7 +512,7 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
         if (psiStatements == null) {
           return null;
         }
-        result = new SwitchStatementBranch(psiStatements);
+        result = new SwitchStatementBranch(psiStatements, statements);
       }
       else {
         PsiExpression returnExpr = returnStmt.getReturnValue();
@@ -525,7 +525,7 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
           if (psiStatements == null) {
             return null;
           }
-          result = new SwitchStatementBranch(psiStatements);
+          result = new SwitchStatementBranch(psiStatements, statements);
         }
         hasReturningBranch = true;
       }
@@ -909,23 +909,40 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
   }
 
   private static final class SwitchStatementBranch implements SwitchRuleResult {
-    final PsiStatement[] myResultStatements;
 
-    private SwitchStatementBranch(PsiStatement[] resultStatements) {
+    private final @Nullable PsiStatement @Nullable [] myResultStatements;
+
+    private final @Nullable PsiStatement @Nullable [] myOriginalResultStatements;
+
+    private SwitchStatementBranch(@Nullable PsiStatement @Nullable[] resultStatements) {
       myResultStatements = resultStatements;
+      myOriginalResultStatements = null;
+    }
+
+    private SwitchStatementBranch(@Nullable PsiStatement @Nullable[] resultStatements, @Nullable PsiStatement @Nullable[] originalResultStatements) {
+      myResultStatements = resultStatements;
+      myOriginalResultStatements = originalResultStatements;
     }
 
     @Override
     public String generate(CommentTracker ct, SwitchBranch branch) {
-      if (myResultStatements.length == 1) {
-        PsiStatement first = myResultStatements[0];
-        if (first instanceof PsiExpressionStatement || first instanceof PsiBlockStatement || first instanceof PsiThrowStatement) return ct.textWithComments(myResultStatements[0]) + "\n";
+      @Nullable PsiStatement @Nullable [] resultStatements = myResultStatements;
+      if(resultStatements == null) return "";
+      if (resultStatements.length == 1) {
+        PsiStatement first = resultStatements[0];
+        if (first instanceof PsiExpressionStatement || first instanceof PsiBlockStatement || first instanceof PsiThrowStatement) {
+          return ct.textWithComments(resultStatements[0]) + "\n";
+        }
       }
       StringBuilder sb = new StringBuilder("{");
-      for (int i = 0, length = myResultStatements.length; i < length; i++) {
-        PsiStatement element = myResultStatements[i];
+      for (int i = 0, length = resultStatements.length; i < length; i++) {
+        PsiStatement element = resultStatements[i];
+        if (element == null) continue;
         if (i == 0) {
-          PsiElement current = element.getPrevSibling();
+          PsiElement current = getElementForComments(element, i);
+          if(current != null) {
+            current = current.getPrevSibling();
+          }
           while (current instanceof PsiWhiteSpace || current instanceof PsiComment) {
             current = current.getPrevSibling();
           }
@@ -933,19 +950,16 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
         }
         sb.append(ct.text(element));
         if (i + 1 < length) {
-          addWhiteSpaceAndComments(element, sb, ct);
+          PsiElement current = getElementForComments(element, i);
+          addWhiteSpaceAndComments(current, sb, ct);
         }
         if (element.getNextSibling() == null && element.getLastChild() instanceof PsiComment comment &&
             comment.getTokenType() == JavaTokenType.END_OF_LINE_COMMENT) {
-          if (sb.charAt(sb.length() - 1) != '\n') {
-            sb.append("\n");
-          }
+          addNewLine(sb);
         }
       }
       addCommentsUntilNextLabel(ct, branch, sb);
-      if (sb.charAt(sb.length() - 1) != '\n') {
-        sb.append("\n");
-      }
+      addNewLine(sb);
       sb.append("}");
       for (PsiElement element : branch.myUsedElements) {
         ct.markUnchanged(element);
@@ -953,6 +967,25 @@ public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLoc
       return sb.toString();
     }
 
+    private static void addNewLine(@NotNull StringBuilder sb) {
+      String string = sb.toString();
+      String trimmed = string.trim();
+      if(!string.substring(trimmed.length()).contains("\n")) {
+        sb.append("\n");
+      }
+    }
+
+    @Nullable
+    private PsiElement getElementForComments(@Nullable PsiStatement element, int i) {
+      PsiElement current = element;
+      if (myOriginalResultStatements != null &&
+          myOriginalResultStatements.length > i &&
+          myOriginalResultStatements[i] != null) {
+        current = myOriginalResultStatements[i];
+      }
+      return current;
+    }
+
     private static void addWhiteSpaceAndComments(@Nullable PsiElement element, @NotNull StringBuilder sb, CommentTracker ct) {
       if (element == null) {
         return;
diff --git a/java/java-tests/testData/inspection/switchExpressionMigration/afterCommentsInTheMiddle.java b/java/java-tests/testData/inspection/switchExpressionMigration/afterCommentsInTheMiddle.java
new file mode 100644
index 000000000000..c928a562ef9d
--- /dev/null
+++ b/java/java-tests/testData/inspection/switchExpressionMigration/afterCommentsInTheMiddle.java
@@ -0,0 +1,18 @@
+// "Fix all 'Statement can be replaced with enhanced 'switch'' problems in file" "true"
+import java.util.*;
+
+class CommentsInTheMiddle {
+  public static String m(int x) {
+      return switch (x) {
+          case 0 ->
+              // first strange comment
+                  "foo";
+          default -> {
+              // some strange comment
+              System.out.print("bar");
+              // another strange comment
+              yield "bar";
+          }
+      };
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/inspection/switchExpressionMigration/beforeCommentsInTheMiddle.java b/java/java-tests/testData/inspection/switchExpressionMigration/beforeCommentsInTheMiddle.java
new file mode 100644
index 000000000000..286f23b854c6
--- /dev/null
+++ b/java/java-tests/testData/inspection/switchExpressionMigration/beforeCommentsInTheMiddle.java
@@ -0,0 +1,17 @@
+// "Fix all 'Statement can be replaced with enhanced 'switch'' problems in file" "true"
+import java.util.*;
+
+class CommentsInTheMiddle {
+  public static String m(int x) {
+    switc<caret>h(x) {
+      case 0:
+        // first strange comment
+        return "foo";
+      default:
+        // some strange comment
+        System.out.print("bar");
+        // another strange comment
+        return "bar";
+    }
+  }
+}
\ No newline at end of file