// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.intellij.codeInspection;

import com.intellij.codeInsight.BlockUtils;
import com.intellij.codeInspection.options.OptPane;
import com.intellij.java.JavaBundle;
import com.intellij.modcommand.ModPsiUpdater;
import com.intellij.modcommand.PsiUpdateModCommandQuickFix;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.pom.java.JavaFeature;
import com.intellij.psi.*;
import com.intellij.psi.controlFlow.*;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.psi.util.PsiUtil;
import com.intellij.psi.util.TypeConversionUtil;
import com.intellij.util.ArrayUtil;
import com.intellij.util.containers.ContainerUtil;
import com.siyeh.ig.psiutils.*;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.Nls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.annotations.PropertyKey;

import java.util.*;

import static com.intellij.codeInspection.options.OptPane.*;
import static com.intellij.util.ObjectUtils.tryCast;
import static com.siyeh.ig.psiutils.SwitchUtils.evaluateSwitchCompleteness;

public final class EnhancedSwitchMigrationInspection extends AbstractBaseJavaLocalInspectionTool {
  @SuppressWarnings("WeakerAccess") public boolean myWarnOnlyOnExpressionConversion = true;
  @SuppressWarnings("WeakerAccess") public int myMaxNumberStatementsForBranch = 2;

  @Override
  public @NotNull OptPane getOptionsPane() {
    return pane(
      checkbox("myWarnOnlyOnExpressionConversion", JavaBundle.message("inspection.switch.expression.migration.warn.only.on.expression")),
      number("myMaxNumberStatementsForBranch", JavaBundle.message("inspection.switch.expression.migration.expression.max.statements"),
             1, 200));
  }

  private static final SwitchConversion[] ourInspections = new SwitchConversion[]{
    (statement, branches, isExhaustive, maxNumberStatementsForExpression) -> inspectReturningSwitch(statement, branches, isExhaustive,
                                                                                                    maxNumberStatementsForExpression),
    (statement, branches, isExhaustive, maxNumberStatementsForExpression) -> inspectVariableAssigningSwitch(statement, branches,
                                                                                                            isExhaustive,
                                                                                                            maxNumberStatementsForExpression),
    (statement, branches, isExhaustive, maxNumberStatementsForExpression) -> inspectReplacementWithStatement(statement, branches)
  };

  @Override
  public @NotNull Set<@NotNull JavaFeature> requiredFeatures() {
    return Set.of(JavaFeature.ENHANCED_SWITCH);
  }

  @Override
  public @NotNull PsiElementVisitor buildVisitor(@NotNull ProblemsHolder holder, boolean isOnTheFly) {
    return new JavaElementVisitor() {

      @Override
      public void visitSwitchExpression(@NotNull PsiSwitchExpression expression) {
        PsiElement switchKeyword = expression.getFirstChild();
        if (switchKeyword == null) {
          return;
        }
        PsiCodeBlock body = expression.getBody();
        if (body == null) return;
        boolean onlyOneYieldAfterLabel = true;
        boolean isOldSwitchWithoutRule = true;
        int statementAfterLabelCount = 0;
        PsiStatement[] statements = body.getStatements();
        if (statements.length == 0) return;
        for (PsiStatement statement : statements) {
          if (statement instanceof PsiSwitchLabeledRuleStatement) {
            isOldSwitchWithoutRule = false;
            break;
          }
          if (statement instanceof PsiSwitchLabelStatement) {
            statementAfterLabelCount = 0;
          }
          else {
            statementAfterLabelCount++;
            if (!(statement instanceof PsiYieldStatement || statement instanceof PsiThrowStatement)) {
              onlyOneYieldAfterLabel = false;
            }
            else {
              if (statementAfterLabelCount > 1) {
                onlyOneYieldAfterLabel = false;
              }
            }
          }
        }
        if (!isOldSwitchWithoutRule) {
          return;
        }
        ProblemHighlightType warningType = ProblemHighlightType.GENERIC_ERROR_OR_WARNING;
        if (!onlyOneYieldAfterLabel) {
          warningType = ProblemHighlightType.INFORMATION;
        }
        holder.registerProblem(switchKeyword,
                               JavaBundle.message("inspection.switch.expression.migration.inspection.switch.expression.description"),
                               warningType, new ReplaceExpressionWithEnhancedSwitchExpressionFix());
      }

      @Override
      public void visitSwitchStatement(@NotNull PsiSwitchStatement statement) {
        PsiElement switchKeyword = statement.getFirstChild();
        if (switchKeyword == null) {
          return;
        }
        List<SwitchReplacer> replacers = findSwitchReplacers(statement, myMaxNumberStatementsForBranch);
        if (replacers.isEmpty()) return;
        Optional<SwitchReplacer> replacerWithWarningLevel = replacers.stream().filter(t -> isWarningLevel(t)).findFirst();
        if (replacerWithWarningLevel.isPresent()) {
          SwitchReplacer replacer = replacerWithWarningLevel.get();
          List<LocalQuickFix> fixes = new ArrayList<>();
          fixes.add(new ReplaceWithSwitchExpressionFix(replacer.getType()));
          if (!myWarnOnlyOnExpressionConversion && replacer.getType() == ReplacementType.Statement) {
            fixes.add(LocalQuickFix.from(new UpdateInspectionOptionFix(
              EnhancedSwitchMigrationInspection.this, "myWarnOnlyOnExpressionConversion",
              JavaBundle.message("inspection.switch.expression.migration.warn.only.on.expression"),
              true)));
          }
          if (replacer.getType() == ReplacementType.Expression &&
              replacer.getMaxNumberStatementsInBranch() != null &&
              replacer.getMaxNumberStatementsInBranch() > 1) {
            int newMaxValue = replacer.getMaxNumberStatementsInBranch() - 1;
            fixes.add(LocalQuickFix.from(new UpdateInspectionOptionFix(
              EnhancedSwitchMigrationInspection.this, "myMaxNumberStatementsForBranch",
              JavaBundle.message("inspection.switch.expression.migration.option.expression.max.statements", newMaxValue),
              newMaxValue)));
          }
          holder.registerProblem(switchKeyword, JavaBundle.message("inspection.switch.expression.migration.inspection.switch.description"),
                                 ProblemHighlightType.GENERIC_ERROR_OR_WARNING, fixes.toArray(LocalQuickFix.EMPTY_ARRAY));
          replacers.remove(replacer);
        }
        if (!holder.isOnTheFly()) {
          return;
        }
        if (replacers.isEmpty()) {
          return;
        }
        List<LocalQuickFix> fixes = ContainerUtil.map(replacers, replacer -> new ReplaceWithSwitchExpressionFix(replacer.getType()));
        holder.registerProblem(switchKeyword, JavaBundle.message("inspection.switch.expression.migration.inspection.switch.description"),
                               ProblemHighlightType.INFORMATION, fixes.toArray(LocalQuickFix.EMPTY_ARRAY));
      }

      private boolean isWarningLevel(@NotNull SwitchReplacer replacer) {
        if (replacer.isInformLevel()) {
          return false;
        }
        return !(myWarnOnlyOnExpressionConversion && replacer.getType() == ReplacementType.Statement);
      }
    };
  }

  private static List<SwitchReplacer> runInspections(@NotNull PsiStatement statement,
                                                     boolean isExhaustive,
                                                     @NotNull List<OldSwitchStatementBranch> branches,
                                                     int maxNumberStatementsForExpression) {
    List<SwitchReplacer> replacers = new ArrayList<>();
    for (SwitchConversion inspection : ourInspections) {
      SwitchReplacer replacer = inspection.suggestReplacer(statement, branches, isExhaustive, maxNumberStatementsForExpression);
      if (replacer != null) {
        replacers.add(replacer);
      }
    }
    return replacers;
  }

  private static OldSwitchStatementBranch addBranch(List<? super OldSwitchStatementBranch> branches,
                                                    PsiStatement[] statements,
                                                    int unmatchedCaseIndex,
                                                    int endIndexExcl,
                                                    boolean isFallthrough, PsiBreakStatement current) {
    PsiSwitchLabelStatement labelStatement = (PsiSwitchLabelStatement)statements[unmatchedCaseIndex];
    PsiStatement[] branchStatements = Arrays.copyOfRange(statements, unmatchedCaseIndex + 1, endIndexExcl);
    OldSwitchStatementBranch branch = new OldSwitchStatementBranch(isFallthrough, branchStatements, labelStatement, current);
    branches.add(branch);
    return branch;
  }

  private static @Nullable List<OldSwitchStatementBranch> extractBranches(@NotNull PsiCodeBlock body,
                                                                          PsiSwitchStatement switchStatement) {
    List<OldSwitchStatementBranch> branches = new ArrayList<>();
    PsiStatement[] statements = body.getStatements();
    int unmatchedCaseIndex = -1;
    OldSwitchStatementBranch previousBranch = null;
    for (int i = 0, length = statements.length; i < length; i++) {
      PsiStatement current = statements[i];
      if (current instanceof PsiSwitchLabelStatement) {
        if (unmatchedCaseIndex != -1) {
          boolean isFallthrough = i != 0 && ControlFlowUtils.statementMayCompleteNormally(statements[i - 1]);
          OldSwitchStatementBranch newBranch = addBranch(branches, statements, unmatchedCaseIndex, i, isFallthrough, null);
          newBranch.myPreviousSwitchBranch = previousBranch;
          previousBranch = newBranch;
        }
        unmatchedCaseIndex = i;
      }
      else if (current instanceof PsiBreakStatement) {
        if (((PsiBreakStatement)current).findExitedStatement() != switchStatement) return null;
        if (unmatchedCaseIndex == -1) return null;
        OldSwitchStatementBranch newBranch = addBranch(branches, statements, unmatchedCaseIndex, i, false, (PsiBreakStatement)current);
        newBranch.myPreviousSwitchBranch = previousBranch;
        previousBranch = newBranch;
        unmatchedCaseIndex = -1;
      }
      else if (current instanceof PsiSwitchLabeledRuleStatement) {
        return null;
      }
    }
    // tail
    if (unmatchedCaseIndex != -1) {
      OldSwitchStatementBranch branch = addBranch(branches, statements, unmatchedCaseIndex, statements.length, false, null);
      branch.myPreviousSwitchBranch = previousBranch;
    }
    return branches;
  }

  /**
   * Before using this method, make sure you are using a correct version of Java.
   */
  public static @Nullable SwitchReplacer findSwitchReplacer(@NotNull PsiSwitchStatement switchStatement) {
    List<SwitchReplacer> replacers = findSwitchReplacers(switchStatement, 1);
    for (SwitchReplacer replacer : replacers) {
      if (!replacer.isInformLevel()) {
        return replacer;
      }
    }
    return null;
  }

  private static @NotNull List<SwitchReplacer> findSwitchReplacers(@NotNull PsiSwitchStatement switchStatement,
                                                                   int maxNumberStatementsForExpression) {
    PsiExpression expression = switchStatement.getExpression();
    if (expression == null) return List.of();
    PsiCodeBlock body = switchStatement.getBody();
    if (body == null) return List.of();
    List<OldSwitchStatementBranch> branches = extractBranches(body, switchStatement);
    if (branches == null || branches.isEmpty()) return List.of();
    boolean isExhaustive = isExhaustiveSwitch(branches, switchStatement);
    return runInspections(switchStatement, isExhaustive, branches, maxNumberStatementsForExpression);
  }

  private static @Nullable PsiSwitchBlock generateEnhancedSwitch(@NotNull PsiStatement statementToReplace,
                                                                 List<SwitchBranch> newBranches,
                                                                 CommentTracker ct,
                                                                 boolean isExpr) {
    PsiElementFactory factory = JavaPsiFacade.getElementFactory(statementToReplace.getProject());
    if (!(statementToReplace instanceof PsiSwitchStatement)) return null;
    PsiCodeBlock body = ((PsiSwitchStatement)statementToReplace).getBody();
    if (body == null) return null;

    StringBuilder sb = new StringBuilder();
    for (PsiElement e = statementToReplace.getFirstChild(); e != null && e != body; e = e.getNextSibling()) {
      sb.append(ct.text(e));
    }
    PsiJavaToken lBrace = body.getLBrace();
    sb.append(lBrace != null ? ct.textWithComments(lBrace) : "{");
    for (SwitchBranch newBranch : newBranches) {
      sb.append(newBranch.generate(ct));
    }
    PsiJavaToken rBrace = body.getRBrace();
    sb.append(rBrace != null ? ct.textWithComments(rBrace) : "}");
    PsiSwitchBlock switchBlock;
    if (isExpr) {
      switchBlock = (PsiSwitchBlock)factory.createExpressionFromText(sb.toString(), statementToReplace);
    }
    else {
      switchBlock = (PsiSwitchBlock)factory.createStatementFromText(sb.toString(), statementToReplace);
    }
    StreamEx.ofTree((PsiElement)switchBlock, block -> Arrays.stream(block.getChildren()))
      .select(PsiBreakStatement.class)
      .filter(breakStmt -> ControlFlowUtils.statementCompletesWithStatement(switchBlock, breakStmt) &&
                           breakStmt.findExitedStatement() == switchBlock)
      .forEach(statement -> new CommentTracker().delete(statement));
    return switchBlock;
  }

  private static boolean isExhaustiveSwitch(List<OldSwitchStatementBranch> branches, PsiSwitchStatement switchStatement) {
    for (OldSwitchStatementBranch branch : branches) {
      if (branch.isDefault()) return true;
      if (existsDefaultLabelElement(branch.myLabelStatement)) return true;
    }
    SwitchUtils.SwitchExhaustivenessState completenessResult = evaluateSwitchCompleteness(switchStatement, true);
    return completenessResult == SwitchUtils.SwitchExhaustivenessState.EXHAUSTIVE_CAN_ADD_DEFAULT ||
           completenessResult == SwitchUtils.SwitchExhaustivenessState.EXHAUSTIVE_NO_DEFAULT;
  }

  private static boolean isConvertibleBranch(@NotNull OldSwitchStatementBranch branch, boolean hasNext) {
    int length = branch.getStatements().length;
    if (length == 0) return (branch.isFallthrough() && hasNext) || (!branch.isFallthrough() && branch.isDefault());
    return !branch.isFallthrough();
  }

  private enum ReplacementType {
    Expression("inspection.replace.with.switch.expression.fix.name"),
    Statement("inspection.replace.with.enhanced.switch.statement.fix.name");
    private final @PropertyKey(resourceBundle = JavaBundle.BUNDLE) String key;

    ReplacementType(@PropertyKey(resourceBundle = JavaBundle.BUNDLE) String key) {
      this.key = key;
    }

    @Nls
    String getFixName() {
      return JavaBundle.message(key);
    }
  }

  public interface SwitchReplacer {
    void replace(@NotNull PsiStatement switchStatement);

    ReplacementType getType();

    boolean isInformLevel();

    //if null, it is not applicable
    @Nullable
    Integer getMaxNumberStatementsInBranch();
  }

  private interface SwitchConversion {
    @Nullable
    SwitchReplacer suggestReplacer(@NotNull PsiStatement statement,
                                   @NotNull List<OldSwitchStatementBranch> branches,
                                   boolean isExhaustive,
                                   int maxNumberStatementsForExpression);
  }

  //Right part of switch rule (case labels -> result)
  private interface SwitchRuleResult {
    String generate(CommentTracker ct, SwitchBranch branch);
  }

  private static class ReplaceExpressionWithEnhancedSwitchExpressionFix extends PsiUpdateModCommandQuickFix {

    @Override
    protected void applyFix(@NotNull Project project, @NotNull PsiElement element, @NotNull ModPsiUpdater updater) {
      StringBuilder builder = new StringBuilder();
      if (!(element.getParent() instanceof PsiSwitchExpression switchExpression)) return;
      for (@NotNull PsiElement switchExpressionChild : switchExpression.getChildren()) {
        if (switchExpressionChild instanceof PsiCodeBlock codeBlock) {
          boolean previousHasBodyStatement = true;
          for (@NotNull PsiElement codeBlockChildren : codeBlock.getChildren()) {
            if (codeBlockChildren instanceof PsiWhiteSpace && builder.charAt(builder.length() - 1) == ',') {
              continue;
            }
            if (codeBlockChildren instanceof PsiSwitchLabelStatement switchLabelStatement) {
              boolean nextIsBodyStatement = checkNextIsStatement(codeBlockChildren);
              for (@NotNull PsiElement labelStatementChild : switchLabelStatement.getChildren()) {
                if (!previousHasBodyStatement &&
                    labelStatementChild instanceof PsiJavaToken javaToken &&
                    javaToken.textMatches("case")) {
                  continue;
                }
                if (labelStatementChild instanceof PsiJavaToken javaToken && javaToken.textMatches(":")) {
                  if (nextIsBodyStatement) {
                    builder.append("->"); //replace ':' with '->'
                  }
                  else {
                    //next label
                    builder.append(",");
                  }
                }
                else {
                  builder.append(labelStatementChild.getText());
                }
              }
              PsiElement nextOfSwitchLabelStatement = PsiTreeUtil.skipWhitespacesAndCommentsForward(switchLabelStatement);
              if (nextIsBodyStatement && !(nextOfSwitchLabelStatement instanceof PsiBlockStatement) &&
                  findOneYieldOrThrowStatement(PsiTreeUtil.skipWhitespacesAndCommentsForward(codeBlockChildren)) == null) {
                builder.append("{"); //wrap multiline rule into '{}'
              }
              previousHasBodyStatement = nextIsBodyStatement;
            }
            else {
              PsiStatement yieldStatement = findOneYieldOrThrowStatement(codeBlockChildren);
              if (yieldStatement != null) {
                for (@NotNull PsiElement yieldStatementChild : yieldStatement.getChildren()) {
                  if (!(yieldStatementChild instanceof PsiJavaToken javaToken && javaToken.textMatches("yield"))) {
                    builder.append(yieldStatementChild.getText()); // skip 'yield' for one-line rule
                  }
                }
              }
              else {
                builder.append(codeBlockChildren.getText());
                PsiElement nextOfCodeBlockChildren = PsiTreeUtil.skipWhitespacesAndCommentsForward(codeBlockChildren);
                if (!(codeBlockChildren instanceof PsiComment || codeBlockChildren instanceof PsiWhiteSpace) &&
                    !(PsiTreeUtil.skipWhitespacesAndCommentsBackward(codeBlockChildren) instanceof PsiJavaToken) &&
                    findOneYieldOrThrowStatement(PsiTreeUtil.skipWhitespacesAndCommentsBackward(codeBlockChildren)) == null &&
                    !(codeBlockChildren instanceof PsiJavaToken) &&
                    (nextOfCodeBlockChildren instanceof PsiSwitchLabelStatement ||
                     nextOfCodeBlockChildren instanceof PsiJavaToken javaToken && javaToken.textMatches("}"))) {
                  builder.append("\n}"); //wrap multiline rule into '{}'
                }
              }
            }
          }
        }
        else {
          builder.append(switchExpressionChild.getText());
        }
      }
      PsiElementFactory factory = JavaPsiFacade.getElementFactory(element.getProject());
      PsiExpression newSwitchExpression = factory.createExpressionFromText(builder.toString(), element);
      switchExpression.replace(newSwitchExpression);
    }

    private static boolean checkNextIsStatement(@Nullable PsiElement statement) {
      PsiElement forward = PsiTreeUtil.skipWhitespacesAndCommentsForward(statement);
      return statement instanceof PsiSwitchLabelStatement &&
             !(forward instanceof PsiSwitchLabelStatement) &&
             forward instanceof PsiStatement;
    }

    private static @Nullable PsiStatement findOneYieldOrThrowStatement(@Nullable PsiElement switchBlockChild) {
      if (switchBlockChild == null) return null;
      boolean isOldOrThrow = switchBlockChild instanceof PsiYieldStatement ||
                             switchBlockChild instanceof PsiThrowStatement;
      if (!isOldOrThrow) {
        return null;
      }
      boolean hasSwitchLabelBefore = PsiTreeUtil.skipWhitespacesAndCommentsBackward(switchBlockChild) instanceof PsiSwitchLabelStatement;
      if (!hasSwitchLabelBefore) {
        return null;
      }
      PsiElement nextElement = PsiTreeUtil.skipWhitespacesAndCommentsForward(switchBlockChild);
      boolean nextElementIsSwitchLabel = nextElement instanceof PsiSwitchLabelStatement;
      boolean isClosingBrace = nextElement instanceof PsiJavaToken javaToken && javaToken.textMatches("}");
      if (!(nextElementIsSwitchLabel || isClosingBrace)) {
        return null;
      }
      return (PsiStatement)switchBlockChild;
    }

    @Override
    public @NotNull String getFamilyName() {
      return JavaBundle.message("inspection.replace.with.switch.rule.expression.fix.family.name");
    }
  }

  private static class ReplaceWithSwitchExpressionFix extends PsiUpdateModCommandQuickFix {
    private final ReplacementType myReplacementType;

    ReplaceWithSwitchExpressionFix(ReplacementType replacementType) { myReplacementType = replacementType; }

    @Override
    public @Nls(capitalization = Nls.Capitalization.Sentence) @NotNull String getFamilyName() {
      return JavaBundle.message("inspection.replace.with.switch.expression.fix.family.name");
    }

    @Override
    public @NotNull String getName() {
      return myReplacementType.getFixName();
    }

    @Override
    protected void applyFix(@NotNull Project project, @NotNull PsiElement element, @NotNull ModPsiUpdater updater) {
      PsiSwitchStatement statement = PsiTreeUtil.getParentOfType(element, PsiSwitchStatement.class);
      if (statement == null) return;
      SwitchReplacer replacer =
        ContainerUtil.find(findSwitchReplacers(statement, Integer.MAX_VALUE), t -> t.getType() == myReplacementType);
      if (replacer == null) return;
      replacer.replace(statement);
    }
  }

  /**
   * This method is used to rearrange branches.
   * Now it can change the order of branches or divide some branches into several separate branches.
   * For example, a null branch will be extracted from others.
   *
   * @return rearranged branches
   */
  private static @NotNull List<SwitchBranch> rearrangeBranches(@NotNull List<SwitchBranch> branches,
                                                               @NotNull PsiElement context) {
    if (branches.isEmpty()) {
      return branches;
    }
    if (!PsiUtil.isAvailable(JavaFeature.PATTERNS_IN_SWITCH, context)) {
      return branches;
    }
    List<SwitchBranch> result = new ArrayList<>();
    for (SwitchBranch branch : branches) {
      if (branch.myIsDefault) {
        result.add(branch);
        continue;
      }
      List<? extends PsiCaseLabelElement> caseExpressions = branch.myCaseExpressions;
      if (caseExpressions == null || caseExpressions.size() <= 1) {
        result.add(branch);
        continue;
      }
      rearrangeCases(branch, caseExpressions, result);
    }
    return result;
  }

  /**
   * This method is used to rearrange case elements.
   * Method is used by {@link #rearrangeBranches(List, PsiElement)}.
   *
   * @param result - container, where sorted cases will be added
   */
  private static void rearrangeCases(@NotNull SwitchBranch branch,
                                     @NotNull List<? extends PsiCaseLabelElement> caseExpressions,
                                     @NotNull List<SwitchBranch> result) {
    List<PsiCaseLabelElement> previousExpressions = new ArrayList<>();
    for (PsiCaseLabelElement expression : caseExpressions) {
      PsiCaseLabelElement external = findNullLabel(List.of(expression));
      if (external == null && expression instanceof PsiPattern) {
        external = expression;
      }
      if (external == null) {
        previousExpressions.add(expression);
      }
      else {
        if (!previousExpressions.isEmpty()) {
          SwitchBranch otherBranch = branch.withLabels(new ArrayList<>(previousExpressions));
          result.add(otherBranch);
          previousExpressions.clear();
        }
        SwitchBranch specialBranch = branch.withLabels(List.of(external));
        result.add(specialBranch);
      }
    }
    if (!previousExpressions.isEmpty()) {
      SwitchBranch otherBranch = branch.withLabels(previousExpressions);
      result.add(otherBranch);
    }
  }

  private static @Nullable PsiCaseLabelElement findNullLabel(@NotNull List<? extends PsiCaseLabelElement> expressions) {
    return ContainerUtil.find(expressions,
                              label -> label instanceof PsiExpression literal && TypeConversionUtil.isNullType(literal.getType()));
  }

  private static final class ReturningSwitchReplacer implements SwitchReplacer {
    final @NotNull PsiStatement myStatement;
    final List<SwitchBranch> myNewBranches;
    final @Nullable PsiReturnStatement myReturnToDelete;
    final @Nullable PsiThrowStatement myThrowStatementToDelete;
    private final @NotNull List<? extends PsiStatement> myStatementsToDelete;
    private final boolean myIsInfo;
    private final int myMaxNumberStatementsInBranch;

    private ReturningSwitchReplacer(@NotNull PsiStatement statement,
                                    @NotNull List<SwitchBranch> newBranches,
                                    @Nullable PsiReturnStatement returnToDelete,
                                    @Nullable PsiThrowStatement throwToDelete,
                                    @NotNull List<? extends PsiStatement> statementsToDelete,
                                    boolean isInfo,
                                    int maxNumberStatementsInBranch) {
      myStatement = statement;
      myNewBranches = rearrangeBranches(newBranches, statement);
      myReturnToDelete = returnToDelete;
      myThrowStatementToDelete = throwToDelete;
      myStatementsToDelete = statementsToDelete;
      myIsInfo = isInfo;
      myMaxNumberStatementsInBranch = maxNumberStatementsInBranch;
    }

    @Override
    public Integer getMaxNumberStatementsInBranch() {
      return myMaxNumberStatementsInBranch;
    }

    @Override
    public boolean isInformLevel() {
      return myIsInfo;
    }

    @Override
    public void replace(@NotNull PsiStatement statement) {
      CommentTracker commentTracker = new CommentTracker();
      PsiSwitchBlock switchBlock = generateEnhancedSwitch(statement, myNewBranches, commentTracker, true);
      if (switchBlock == null) return;

      if (myReturnToDelete != null) {
        CommentTracker ct = new CommentTracker();
        commentTracker.markUnchanged(myReturnToDelete.getReturnValue());
        ct.delete(myReturnToDelete);
      }
      if (myThrowStatementToDelete != null) {
        CommentTracker ct = new CommentTracker();
        commentTracker.markUnchanged(myThrowStatementToDelete.getException());
        ct.delete(myThrowStatementToDelete);
      }
      for (PsiStatement toDelete : myStatementsToDelete) {
        commentTracker.delete(toDelete);
      }
      PsiElementFactory factory = JavaPsiFacade.getElementFactory(statement.getProject());
      PsiStatement returnStatement = factory.createStatementFromText("return " + switchBlock.getText() + ";", switchBlock);
      commentTracker.replaceAndRestoreComments(statement, returnStatement);
    }

    @Override
    public ReplacementType getType() {
      return ReplacementType.Expression;
    }
  }

  /**
   * <pre>
   * switch (n) {
   *   case 1:
   *     return "a";
   *   case 2:
   *     return "b";
   *   default:
   *     return "?";
   * }
   * </pre>
   */

  private static @Nullable SwitchReplacer inspectReturningSwitch(@NotNull PsiStatement statement,
                                                                 @NotNull List<OldSwitchStatementBranch> branches,
                                                                 boolean isExhaustive, int maxNumberStatementsForExpression) {
    PsiReturnStatement returnAfterSwitch =
      tryCast(PsiTreeUtil.getNextSiblingOfType(statement, PsiStatement.class), PsiReturnStatement.class);
    PsiThrowStatement throwAfterSwitch =
      tryCast(PsiTreeUtil.getNextSiblingOfType(statement, PsiStatement.class), PsiThrowStatement.class);
    if (returnAfterSwitch == null && throwAfterSwitch == null && !isExhaustive) return null;
    List<SwitchBranch> newBranches = new ArrayList<>();
    boolean hasReturningBranch = false;
    boolean isInfo = false;
    int maxLines = 0;
    for (int i = 0, size = branches.size(); i < size; i++) {
      OldSwitchStatementBranch branch = branches.get(i);
      if (!isConvertibleBranch(branch, i != size - 1)) return null;
      if (branch.isFallthrough()) continue;
      PsiStatement[] statements = branch.getStatements();
      if (statements.length == 1 && statements[0] instanceof PsiBlockStatement psiCodeBlock) {
        statements = psiCodeBlock.getCodeBlock().getStatements();
      }
      if (statements.length == 0) {
        if ((i == branches.size() - 1) || branch.isDefault()) {
          if (returnAfterSwitch != null) {
            statements = new PsiStatement[]{returnAfterSwitch};
          }
          else if (throwAfterSwitch != null) {
            statements = new PsiStatement[]{throwAfterSwitch};
          }
          else {
            return null;
          }
        }
        else {
          return null;
        }
      }
      if (maxLines < statements.length) {
        maxLines = statements.length;
      }
      if (statements.length > maxNumberStatementsForExpression) {
        isInfo = true;
      }
      int lastIndex = statements.length - 1;
      if (ContainerUtil.exists(statements,
                               st -> !PsiTreeUtil.findChildrenOfAnyType(st, PsiContinueStatement.class, PsiBreakStatement.class,
                                                                        PsiYieldStatement.class).isEmpty())) {
        return null;
      }
      PsiReturnStatement returnStmt = tryCast(statements[lastIndex], PsiReturnStatement.class);
      SwitchRuleResult result;
      if (returnStmt == null) {
        PsiThrowStatement throwStatement = tryCast(statements[lastIndex], PsiThrowStatement.class);
        if (throwStatement == null) return null;
        PsiStatement[] psiStatements = replaceAllReturnWithYield(statements);
        if (psiStatements == null) {
          return null;
        }
        result = new SwitchStatementBranch(psiStatements, statements);
      }
      else {
        PsiExpression returnExpr = returnStmt.getReturnValue();
        if (returnExpr == null) return null;
        if (statements.length == 1) {
          result = new SwitchRuleExpressionResult(returnExpr);
        }
        else {
          PsiStatement[] psiStatements = replaceAllReturnWithYield(statements);
          if (psiStatements == null) {
            return null;
          }
          result = new SwitchStatementBranch(psiStatements, statements);
        }
        hasReturningBranch = true;
      }
      newBranches.add(SwitchBranch.fromOldBranch(branch, result, branch.getUsedElements()));
    }
    if (!hasReturningBranch) return null;
    if (!isExhaustive && returnAfterSwitch != null) {
      PsiExpression returnExpr = returnAfterSwitch.getReturnValue();
      if (returnExpr == null) return null;
      newBranches.add(SwitchBranch.createDefault(new SwitchRuleExpressionResult(returnExpr)));
    }
    if (!isExhaustive && throwAfterSwitch != null) {
      newBranches.add(SwitchBranch.createDefault(new SwitchStatementBranch(new PsiStatement[]{throwAfterSwitch})));
    }
    List<PsiStatement> statementsToDelete = new ArrayList<>();
