// 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.codeInsight;

import com.intellij.codeInsight.completion.*;
import com.intellij.codeInsight.hints.ParameterHintsPass;
import com.intellij.codeInsight.lookup.CommaTailType;
import com.intellij.openapi.components.Service;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.Condition;
import com.intellij.openapi.util.NullableComputable;
import com.intellij.pom.java.JavaFeature;
import com.intellij.pom.java.LanguageLevel;
import com.intellij.psi.*;
import com.intellij.psi.codeStyle.JavaCodeStyleManager;
import com.intellij.psi.codeStyle.VariableKind;
import com.intellij.psi.impl.source.resolve.CompletionParameterTypeInferencePolicy;
import com.intellij.psi.impl.source.resolve.DefaultParameterTypeInferencePolicy;
import com.intellij.psi.impl.source.resolve.ParameterTypeInferencePolicy;
import com.intellij.psi.infos.CandidateInfo;
import com.intellij.psi.infos.MethodCandidateInfo;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.psi.search.PsiShortNamesCache;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.util.*;
import com.intellij.util.ArrayUtil;
import com.intellij.util.JavaPsiConstructorUtil;
import com.intellij.util.ObjectUtils;
import com.intellij.util.containers.ContainerUtil;
import com.intellij.util.containers.Stack;
import com.siyeh.ig.psiutils.TypeUtils;
import com.siyeh.ig.psiutils.VariableAccessUtils;
import com.siyeh.ig.testFrameworks.AssertHint;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.function.Supplier;

@Service(Service.Level.PROJECT)
public final class ExpectedTypesProvider {
  private static final ExpectedTypeInfo VOID_EXPECTED = createInfoImpl(PsiTypes.voidType(), ExpectedTypeInfo.TYPE_OR_SUBTYPE,
                                                                       PsiTypes.voidType(), TailTypes.semicolonType());

  private static final Logger LOG = Logger.getInstance(ExpectedTypesProvider.class);

  public static ExpectedTypesProvider getInstance(@NotNull Project project) {
    return project.getService(ExpectedTypesProvider.class);
  }

  private static final int MAX_COUNT = 50;
  private static final ExpectedClassProvider ourGlobalScopeClassProvider = new ExpectedClassProvider() {
    @Override
    public PsiField @NotNull [] findDeclaredFields(@NotNull PsiManager manager, @NotNull String name) {
      final PsiShortNamesCache cache = PsiShortNamesCache.getInstance(manager.getProject());
      GlobalSearchScope scope = GlobalSearchScope.allScope(manager.getProject());
      return cache.getFieldsByName(name, scope);
    }

    @Override
    public PsiMethod @NotNull [] findDeclaredMethods(@NotNull PsiManager manager, @NotNull String name) {
      Project project = manager.getProject();
      final PsiShortNamesCache cache = PsiShortNamesCache.getInstance(project);
      GlobalSearchScope sources = GlobalSearchScope.projectScope(project);
      GlobalSearchScope libraries = GlobalSearchScope.notScope(sources);
      PsiMethod[] sourceMethods = cache.getMethodsByNameIfNotMoreThan(name, sources, MAX_COUNT);
      if (sourceMethods.length >= MAX_COUNT) return sourceMethods;
      PsiMethod[] libraryMethods = cache.getMethodsByNameIfNotMoreThan(name, libraries, MAX_COUNT-sourceMethods.length);
      return ArrayUtil.mergeArrays(sourceMethods, libraryMethods);
    }
  };

  public static @NotNull ExpectedTypeInfo createInfo(@NotNull  PsiType type, @ExpectedTypeInfo.Type int kind, PsiType defaultType, @NotNull TailType tailType) {
    return createInfoImpl(type, kind, defaultType, tailType);
  }

  private static @NotNull ExpectedTypeInfoImpl createInfoImpl(@NotNull PsiType type, PsiType defaultType) {
    return createInfoImpl(type, ExpectedTypeInfo.TYPE_OR_SUBTYPE, defaultType, TailTypes.noneType());
  }

  private static @NotNull ExpectedTypeInfoImpl createInfoImpl(@NotNull PsiType type, @ExpectedTypeInfo.Type int kind, PsiType defaultType, @NotNull TailType tailType) {
    return new ExpectedTypeInfoImpl(type, kind, defaultType, tailType, null, ExpectedTypeInfoImpl.NULL);
  }

  private static @NotNull ExpectedTypeInfoImpl createInfoImpl(@NotNull PsiType type,
                                                              @ExpectedTypeInfo.Type int kind,
                                                              PsiType defaultType,
                                                              @NotNull TailType tailType,
                                                              PsiMethod calledMethod,
                                                              Supplier<String> expectedName) {
    return new ExpectedTypeInfoImpl(type, kind, defaultType, tailType, calledMethod, expectedName);
  }

  public static @Nullable ExpectedTypeInfo getSingleExpectedTypeForCompletion(@Nullable PsiExpression expr) {
    ExpectedTypeInfo[] expectedTypes = getExpectedTypes(expr, true, ourGlobalScopeClassProvider, false, false, 1);
    if (expectedTypes.length > 0) {
      return expectedTypes[0];
    }
    return null;
  }

  public static ExpectedTypeInfo @NotNull [] getExpectedTypes(@Nullable PsiExpression expr, boolean forCompletion) {
    return getExpectedTypes(expr, forCompletion, false, false);
  }

  public static ExpectedTypeInfo @NotNull [] getExpectedTypes(@Nullable PsiExpression expr,
                                                              boolean forCompletion,
                                                              boolean voidable,
                                                              boolean usedAfter) {
    return getExpectedTypes(expr, forCompletion, ourGlobalScopeClassProvider, voidable, usedAfter);
  }

  public static ExpectedTypeInfo @NotNull [] getExpectedTypes(@Nullable PsiExpression expr,
                                                              boolean forCompletion,
                                                              ExpectedClassProvider classProvider,
                                                              boolean usedAfter) {
    return getExpectedTypes(expr, forCompletion, classProvider, false, usedAfter);
  }

  public static ExpectedTypeInfo @NotNull [] getExpectedTypes(@Nullable PsiExpression expr,
                                                              boolean forCompletion,
                                                              ExpectedClassProvider classProvider,
                                                              boolean voidable,
                                                              boolean usedAfter) {
    return getExpectedTypes(expr, forCompletion, classProvider, voidable, usedAfter, Integer.MAX_VALUE);
  }

  private static ExpectedTypeInfo @NotNull [] getExpectedTypes(@Nullable PsiExpression expr,
                                                               boolean forCompletion,
                                                               ExpectedClassProvider classProvider,
                                                               boolean voidable,
                                                               boolean usedAfter,
                                                               int maxCandidates) {
    if (expr == null) return ExpectedTypeInfo.EMPTY_ARRAY;
    PsiElement parent = expr.getParent();
    PsiFunctionalExpression functionalExpression = extractFunctionalExpression(expr);
    if (functionalExpression != null) {
      final Collection<? extends PsiType> types = FunctionalInterfaceSuggester.suggestFunctionalInterfaces(functionalExpression);
      if (types.isEmpty()) {
        return ExpectedTypeInfo.EMPTY_ARRAY;
      }
      else {
        final ExpectedTypeInfo[] result = new ExpectedTypeInfo[types.size()];
        int i = 0;
        for (PsiType type : types) {
          result[i++] =
            new ExpectedTypeInfoImpl(type, ExpectedTypeInfo.TYPE_SAME_SHAPED, type, TailTypes.noneType(), null, ExpectedTypeInfoImpl.NULL);
        }
        return result;
      }
    }
    MyParentVisitor visitor = new MyParentVisitor(expr, forCompletion, classProvider, voidable, usedAfter, maxCandidates);
    if (parent != null) {
      parent.accept(visitor);
    }
    ExpectedTypeInfo[] result = visitor.getResult();
    if (forCompletion) {
      return ContainerUtil.map2Array(result, ExpectedTypeInfo.class, i -> ((ExpectedTypeInfoImpl)i).fixUnresolvedTypes(expr));
    }
    return result;
  }

  private static PsiFunctionalExpression extractFunctionalExpression(PsiExpression expr) {
    PsiElement parent = expr.getParent();
    if (expr instanceof PsiFunctionalExpression && parent instanceof PsiExpressionStatement && !(parent.getParent() instanceof PsiSwitchLabeledRuleStatement)) {
      return (PsiFunctionalExpression)expr;
    }
    parent = PsiTreeUtil.skipParentsOfType(expr, PsiParenthesizedExpression.class);
    if (parent instanceof PsiAssignmentExpression &&
        parent.getParent() instanceof PsiExpressionStatement &&
        PsiTreeUtil.isAncestor(((PsiAssignmentExpression)parent).getLExpression(), expr, false)) {
      return ObjectUtils.tryCast(((PsiAssignmentExpression)parent).getRExpression(), PsiFunctionalExpression.class);
    }
    return null;
  }

  public static PsiType @NotNull [] processExpectedTypes(ExpectedTypeInfo @NotNull [] infos,
                                                         @NotNull PsiTypeVisitor<? extends PsiType> visitor, @NotNull Project project) {
    LinkedHashSet<PsiType> set = new LinkedHashSet<>();
    for (ExpectedTypeInfo info : infos) {
      ExpectedTypeInfoImpl infoImpl = (ExpectedTypeInfoImpl)info;

      if (infoImpl.getDefaultType() instanceof PsiClassType) {
        JavaResolveResult result = ((PsiClassType)infoImpl.getDefaultType()).resolveGenerics();
        PsiClass aClass = (PsiClass)result.getElement();
        if (aClass instanceof PsiAnonymousClass) {
          processType(((PsiAnonymousClass)aClass).getBaseClassType(), visitor, set);
          ((PsiAnonymousClass)aClass).getBaseClassType().accept(visitor);
        }
        else {
          processType(infoImpl.getDefaultType(), visitor, set);
        }
      }
      else {
        processType(infoImpl.getDefaultType(), visitor, set);
      }

      if (infoImpl.getKind() == ExpectedTypeInfo.TYPE_OR_SUPERTYPE) {
        processAllSuperTypes(infoImpl.getType(), visitor, project, set, new HashSet<>());
      }
      else if (infoImpl.getKind() == ExpectedTypeInfo.TYPE_OR_SUBTYPE) {
        if (infoImpl.getType() instanceof PsiPrimitiveType) {
          processPrimitiveTypeAndSubtypes((PsiPrimitiveType)infoImpl.getType(), visitor, set);
        }
        //else too expensive to search
      }
    }

    return set.toArray(PsiType.createArray(set.size()));
  }

  private static void processType(@NotNull PsiType type, @NotNull PsiTypeVisitor<? extends PsiType> visitor, @NotNull Set<? super PsiType> typeSet) {
    PsiType accepted = type.accept(visitor);
    if (accepted != null) typeSet.add(accepted);
  }

  private static void processPrimitiveTypeAndSubtypes(@NotNull PsiPrimitiveType type,
                                                      @NotNull PsiTypeVisitor<? extends PsiType> visitor,
                                                      @NotNull Set<? super PsiType> set) {
    if (type.equals(PsiTypes.booleanType()) || type.equals(PsiTypes.voidType()) || type.equals(PsiTypes.nullType())) return;

    for (PsiPrimitiveType primitiveType : PsiTypes.primitiveTypes()) {
      processType(primitiveType, visitor, set);
      if (primitiveType.equals(type)) return;
    }
  }

  public static void processAllSuperTypes(@NotNull PsiType type, @NotNull PsiTypeVisitor<? extends PsiType> visitor, @NotNull Project project, @NotNull Set<? super PsiType> set, @NotNull Set<? super PsiType> visited) {
    if (!visited.add(type)) return;

    if (type instanceof PsiPrimitiveType) {
      if (type.equals(PsiTypes.booleanType()) || type.equals(PsiTypes.voidType()) || type.equals(PsiTypes.nullType())) return;

      List<PsiPrimitiveType> primitiveTypes = PsiTypes.primitiveTypes();
      Stack<PsiType> stack = new Stack<>();
      for (PsiPrimitiveType primitiveType : primitiveTypes.reversed()) {
        if (primitiveType.equals(type)) break;
        //it is already processed before
        if (primitiveType.equals(PsiTypes.booleanType())) continue;
        stack.push(primitiveType);
      }
      while (!stack.empty()) {
        processType(stack.pop(), visitor, set);
      }
    }
    else{
      PsiManager manager = PsiManager.getInstance(project);
      GlobalSearchScope resolveScope = type.getResolveScope();
      if (resolveScope == null) resolveScope = GlobalSearchScope.allScope(project);
      PsiClassType objectType = PsiType.getJavaLangObject(manager, resolveScope);
      processType(objectType, visitor, set);

      if (type instanceof PsiClassType) {
        for (PsiType superType : type.getSuperTypes()) {
          processType(superType, visitor, set);
          processAllSuperTypes(superType, visitor, project, set, visited);
        }
      }
    }
  }

  private static final class MyParentVisitor extends JavaElementVisitor {
    private static final int MAX_VAR_HOPS = 3;
    private PsiExpression myExpr;
    private int myHops = 0;
    private final boolean myForCompletion;
    private final boolean myUsedAfter;
    private final int myMaxCandidates;
    private final ExpectedClassProvider myClassProvider;
    private final boolean myVoidable;
    final List<ExpectedTypeInfo> myResult = new ArrayList<>();
    private static final @NonNls String LENGTH_SYNTHETIC_ARRAY_FIELD = "length";

    private MyParentVisitor(PsiExpression expr,
                            boolean forCompletion,
                            ExpectedClassProvider classProvider,
                            boolean voidable,
                            boolean usedAfter,
                            int maxCandidates) {
      myExpr = expr;
      myForCompletion = forCompletion;
      myClassProvider = classProvider;
      myVoidable = voidable;
      myUsedAfter = usedAfter;
      myMaxCandidates = maxCandidates;
    }

    public ExpectedTypeInfo @NotNull [] getResult() {
      if (myResult.size() > myMaxCandidates) {
        return myResult.subList(0, myMaxCandidates).toArray(ExpectedTypeInfo.EMPTY_ARRAY);
      }
      return myResult.toArray(ExpectedTypeInfo.EMPTY_ARRAY);
    }

    @Override
    public void visitParenthesizedExpression(@NotNull PsiParenthesizedExpression expression) {
      PsiElement parent = expression.getParent();
      if (parent != null) {
        final MyParentVisitor visitor = new MyParentVisitor(expression, myForCompletion, myClassProvider, myVoidable, myUsedAfter,
                                                            myMaxCandidates);
        parent.accept(visitor);
        for (ExpectedTypeInfo info : visitor.myResult) {
          myResult.add(createInfoImpl(info.getType(), info.getKind(), info.getDefaultType(), JavaTailTypes.RPARENTH, info.getCalledMethod(),
                                      ((ExpectedTypeInfoImpl)info)::getExpectedName));
        }
      }
    }

    @Override
    public void visitAnnotationMethod(@NotNull PsiAnnotationMethod method) {
      if (myExpr == method.getDefaultValue()) {
        final PsiType type = method.getReturnType();
        if (type != null) {
          myResult.add(createInfoImpl(type, ExpectedTypeInfo.TYPE_OR_SUBTYPE, type, TailTypes.semicolonType()));
        }
      }
    }

    @Override
    public void visitReferenceExpression(@NotNull PsiReferenceExpression expression) {
      if (myForCompletion) {
        final MyParentVisitor visitor = new MyParentVisitor(expression, true, myClassProvider, myVoidable, myUsedAfter, myMaxCandidates);
        expression.getParent().accept(visitor);
        myResult.addAll(visitor.myResult);
        return;
      }

      String referenceName = expression.getReferenceName();
      if (referenceName != null) {
        final PsiElement parent = expression.getParent();
        if (parent instanceof PsiMethodCallExpression) {
          Collections.addAll(myResult, findClassesWithDeclaredMethod((PsiMethodCallExpression)parent));
        }
        else if (parent instanceof PsiVariable ||
                 parent instanceof PsiExpression) {
          if (LENGTH_SYNTHETIC_ARRAY_FIELD.equals(referenceName)) {
            myResult.addAll(anyArrayType());
          }
          else {
            Collections.addAll(myResult, findClassesWithDeclaredField(expression));
          }
        }
      }
    }

    @Override
    public void visitExpressionStatement(@NotNull PsiExpressionStatement statement) {
      if (statement.getParent() instanceof PsiSwitchLabeledRuleStatement) {
        PsiSwitchBlock block = ((PsiSwitchLabeledRuleStatement)statement.getParent()).getEnclosingSwitchBlock();
        if (block instanceof PsiSwitchExpression) {
          Collections.addAll(myResult, getExpectedTypes((PsiExpression)block, myForCompletion));
          return;
        }
      }
      if (myVoidable) {
        myResult.add(VOID_EXPECTED);
      }
    }

    @Override
    public void visitYieldStatement(@NotNull PsiYieldStatement statement) {
      PsiSwitchExpression expression = statement.findEnclosingExpression();
      if (expression != null) {
        Collections.addAll(myResult, getExpectedTypes(expression, myForCompletion));
      }
    }

    @Override
    public void visitMethodCallExpression(@NotNull PsiMethodCallExpression expression) {
      myExpr = (PsiExpression)myExpr.getParent();
      expression.getParent().accept(this);
    }

    @Override
    public void visitAnnotationArrayInitializer(@NotNull PsiArrayInitializerMemberValue initializer) {
      PsiElement parent = initializer.getParent();
      while (parent instanceof PsiArrayInitializerMemberValue) {
        parent = parent.getParent();
      }
      final PsiType type;
      if (parent instanceof PsiNameValuePair) {
        type = getAnnotationMethodType((PsiNameValuePair)parent);
      }
      else {
        type = ((PsiAnnotationMethod)parent).getReturnType();
      }
      if (type instanceof PsiArrayType) {
        final PsiType componentType = ((PsiArrayType)type).getComponentType();
        myResult.add(createInfoImpl(componentType, componentType));
      }
    }

    @Override
    public void visitNameValuePair(@NotNull PsiNameValuePair pair) {
      final PsiType type = getAnnotationMethodType(pair);
      if (type == null) return;
      if (type instanceof PsiArrayType) {
        PsiType componentType = ((PsiArrayType)type).getComponentType();
        myResult.add(createInfoImpl(componentType, componentType));
      }
      else {
        myResult.add(createInfoImpl(type, type));
      }
    }

    private static @Nullable PsiType getAnnotationMethodType(@NotNull PsiNameValuePair pair) {
      final PsiReference reference = pair.getReference();
      if (reference != null) {
        final PsiElement method = reference.resolve();
        if (method instanceof PsiMethod) {
          return ((PsiMethod)method).getReturnType();
        }
      }
      return null;
    }

    @Override
    public void visitLambdaExpression(@NotNull PsiLambdaExpression lambdaExpression) {
      super.visitLambdaExpression(lambdaExpression);
      final PsiType functionalInterfaceType = lambdaExpression.getFunctionalInterfaceType();
      final PsiMethod scopeMethod = LambdaUtil.getFunctionalInterfaceMethod(functionalInterfaceType);
      if (scopeMethod != null) {
        visitMethodReturnType(scopeMethod, LambdaUtil.getFunctionalInterfaceReturnType(functionalInterfaceType),
                              JavaFrontendCompletionUtil.insertSemicolonAfter(lambdaExpression));
      }
    }

    @Override
    public void visitReturnStatement(@NotNull PsiReturnStatement statement) {
      final PsiMethod method;
      final PsiType type;
      final boolean tailTypeSemicolon;
      final NavigatablePsiElement psiElement = PsiTreeUtil.getParentOfType(statement, PsiLambdaExpression.class, PsiMethod.class);
      if (psiElement instanceof PsiLambdaExpression) {
        final PsiType functionalInterfaceType = ((PsiLambdaExpression)psiElement).getFunctionalInterfaceType();
        method = LambdaUtil.getFunctionalInterfaceMethod(functionalInterfaceType);
        type = LambdaUtil.getFunctionalInterfaceReturnType(functionalInterfaceType);
        tailTypeSemicolon = JavaFrontendCompletionUtil.insertSemicolonAfter((PsiLambdaExpression)psiElement);
      }
      else if (psiElement instanceof PsiMethod) {
        method = (PsiMethod)psiElement;
        type = method.getReturnType();
        tailTypeSemicolon = true;
      } else {
        method = null;
        type = null;
        tailTypeSemicolon = true;
      }
      if (method != null) {
        visitMethodReturnType(method, type, tailTypeSemicolon);
      }

    }

    private void visitMethodReturnType(PsiMethod scopeMethod, PsiType type, boolean tailTypeSemicolon) {
      if (type != null) {
        Supplier<String> expectedName;
        if (PropertyUtilBase.isSimplePropertyAccessor(scopeMethod)) {
          expectedName = () -> PropertyUtilBase.getPropertyName(scopeMethod);
        }
        else {
          expectedName = ExpectedTypeInfoImpl.NULL;
        }

        myResult.add(createInfoImpl(type, ExpectedTypeInfo.TYPE_OR_SUBTYPE, type,
                                    tailTypeSemicolon ? TailTypes.semicolonType() : TailTypes.noneType(), null, expectedName));
      }
    }

    @Override
    public void visitIfStatement(@NotNull PsiIfStatement statement) {
      myResult.add(
        createInfoImpl(PsiTypes.booleanType(), ExpectedTypeInfo.TYPE_STRICTLY, PsiTypes.booleanType(), JavaTailTypes.IF_RPARENTH));
    }

    @Override
    public void visitWhileStatement(@NotNull PsiWhileStatement statement) {
      myResult.add(
        createInfoImpl(PsiTypes.booleanType(), ExpectedTypeInfo.TYPE_STRICTLY, PsiTypes.booleanType(), JavaTailTypes.WHILE_RPARENTH));
    }

    @Override
    public void visitDoWhileStatement(@NotNull PsiDoWhileStatement statement) {
      myResult.add(
        createInfoImpl(PsiTypes.booleanType(), ExpectedTypeInfo.TYPE_STRICTLY, PsiTypes.booleanType(), JavaTailTypes.WHILE_RPARENTH));
    }

    @Override
    public void visitForStatement(@NotNull PsiForStatement statement) {
      if (myExpr.equals(statement.getCondition())) {
        myResult.add(createInfoImpl(PsiTypes.booleanType(), ExpectedTypeInfo.TYPE_STRICTLY, PsiTypes.booleanType(),
                                    TailTypes.semicolonType()));
      }
    }

    @Override
    public void visitAssertStatement(@NotNull PsiAssertStatement statement) {
      if (statement.getAssertDescription() == myExpr) {
        final PsiClassType stringType = PsiType.getJavaLangString(myExpr.getManager(), myExpr.getResolveScope());
        myResult.add(createInfoImpl(stringType, ExpectedTypeInfo.TYPE_STRICTLY, stringType, TailTypes.semicolonType()));
      }
      else {
        myResult.add(createInfoImpl(PsiTypes.booleanType(), ExpectedTypeInfo.TYPE_STRICTLY, PsiTypes.booleanType(),
                                    TailTypes.semicolonType()));
      }
    }

    @Override
    public void visitSwitchLabelStatement(@NotNull PsiSwitchLabelStatement statement) {
      processGuard(statement);
    }

    @Override
    public void visitSwitchLabeledRuleStatement(@NotNull PsiSwitchLabeledRuleStatement statement) {
      processGuard(statement);
    }

    private void processGuard(@NotNull PsiSwitchLabelStatementBase statement) {
      if (!myExpr.equals(statement.getGuardExpression())) {
        return;
      }
      final PsiSwitchBlock switchBlock = statement.getEnclosingSwitchBlock();
      if (switchBlock != null) {
        final TailType caseTail = JavaTailTypes.forSwitchLabel(switchBlock);
        myResult.add(createInfoImpl(PsiTypes.booleanType(), ExpectedTypeInfo.TYPE_STRICTLY, PsiTypes.booleanType(), caseTail));
      }
    }

    @Override
    public void visitForeachStatement(@NotNull PsiForeachStatement statement) {
      if (myExpr.equals(statement.getIteratedValue())) {
        PsiParameter iterationParameter = statement.getIterationParameter();
        PsiType type = iterationParameter.getType();

        if (PsiTypes.nullType().equals(type)) return;

        PsiType arrayType = type.createArrayType();
        myResult.add(createInfoImpl(arrayType, arrayType));

        PsiManager manager = statement.getManager();
        PsiElementFactory factory = JavaPsiFacade.getElementFactory(manager.getProject());
        PsiClass iterableClass =
          JavaPsiFacade.getInstance(manager.getProject()).findClass("java.lang.Iterable", statement.getResolveScope());
        if (iterableClass != null && iterableClass.getTypeParameters().length == 1) {
          Map<PsiTypeParameter, PsiType> map = new HashMap<>();
          map.put(iterableClass.getTypeParameters()[0], PsiWildcardType.createExtends(manager, type));
          PsiType iterableType = factory.createType(iterableClass, factory.createSubstitutor(map));
          myResult.add(createInfoImpl(iterableType, iterableType));
        }
      }
    }

    @Override
    public void visitSwitchStatement(@NotNull PsiSwitchStatement statement) {
      processSwitchBlock(statement);
    }

    @Override
    public void visitSwitchExpression(@NotNull PsiSwitchExpression expression) {
      processSwitchBlock(expression);
    }

    public void processSwitchBlock(@NotNull PsiSwitchBlock statement) {
      if (statement.getExpression() == this.myExpr) {
        List<ExpectedTypeInfo> collectedTypes = collectFromLabels(statement);
        if (!collectedTypes.isEmpty()) {
          myResult.addAll(collectedTypes);
          return;
        }
      }
      myResult.add(createInfoImpl(PsiTypes.intType(), PsiTypes.intType()));
      LanguageLevel level = PsiUtil.getLanguageLevel(statement);
      if (JavaFeature.ENUMS.isSufficient(level)) {
        PsiClassType enumType = TypeUtils.getType(CommonClassNames.JAVA_LANG_ENUM, statement);
        myResult.add(createInfoImpl(enumType, enumType));
      }
      if (JavaFeature.STRING_SWITCH.isSufficient(level)) {
        PsiClassType stringType = TypeUtils.getStringType(statement);
        myResult.add(createInfoImpl(stringType, stringType));
      }
    }

    private static @NotNull List<ExpectedTypeInfo> collectFromLabels(@NotNull PsiSwitchBlock statement) {

      List<PsiType> labeledExpressionTypes = new ArrayList<>();
      List<PsiType> labeledPatternsTypes = new ArrayList<>();

      List<ExpectedTypeInfo> result = new ArrayList<>();

      boolean mustBeReference = false;
      PsiCodeBlock body = statement.getBody();
      if (body == null) {
        return result;
      }
      for (PsiStatement psiStatement : body.getStatements()) {
        if (psiStatement instanceof PsiSwitchLabelStatementBase labelStatement) {
          if (labelStatement.isDefaultCase()) {
            continue;
          }
          PsiCaseLabelElementList labelElementList = labelStatement.getCaseLabelElementList();
          if (labelElementList == null) {
            continue;
          }
          for (PsiCaseLabelElement caseLabelElement : labelElementList.getElements()) {
            if (caseLabelElement instanceof PsiExpression expression) {
              PsiType type = expression.getType();
              if (type == null) {
                continue;
              }
              if (type == PsiTypes.nullType()) {
                mustBeReference = true;
                continue;
              }
              labeledExpressionTypes.add(type);
            }
            if (caseLabelElement instanceof PsiPattern pattern) {
              PsiType type = JavaPsiPatternUtil.getPatternType(pattern);
              if (type == null) {
                continue;
              }
              labeledPatternsTypes.add(type);
            }
          }
        }
      }
      result.addAll(findExpectedTypesForExpressions(labeledExpressionTypes, mustBeReference, statement));
      result.addAll(findExpectedTypesForPatterns(labeledPatternsTypes, statement));
      return result;
    }


    private static @NotNull List<ExpectedTypeInfo> findExpectedTypesForPatterns(@NotNull List<PsiType> expectedTypes,
                                                                                @NotNull PsiSwitchBlock statement) {
      PsiManager manager = statement.getManager();
      if(manager == null) return Collections.emptyList();
      PsiType lub = getLeastUpperBound(expectedTypes, manager);
      if (lub instanceof PsiIntersectionType intersectionType) {
        Condition<PsiType> condition = type ->
          !(type instanceof PsiClassType classType && classType.resolve() instanceof PsiClass aClass &&
            CommonClassNames.JAVA_LANG_RECORD.equals(aClass.getQualifiedName()));
        List<PsiType> listWithoutRecords =
          ContainerUtil.filter(intersectionType.getConjuncts(), condition);
        lub = PsiIntersectionType.createIntersection(listWithoutRecords);
      }
      if (lub == null) return Collections.emptyList();
      return Collections.singletonList(createInfo(lub, ExpectedTypeInfo.TYPE_OR_SUPERTYPE, lub, TailTypes.noneType()));
    }

    private static @Nullable PsiType getLeastUpperBound(@NotNull List<PsiType> types, @NotNull PsiManager manager) {
      if (types.isEmpty()) return null;
      Iterator<PsiType> iterator = types.iterator();
      PsiType accumulator = iterator.next();
      while (iterator.hasNext()) {
        PsiType type = iterator.next();
        accumulator = GenericsUtil.getLeastUpperBound(accumulator, type, manager);
        if (accumulator == null) return null;
      }
      return accumulator;
    }

    private static @NotNull List<ExpectedTypeInfo> findExpectedTypesForExpressions(@NotNull List<PsiType> expectedTypes,
                                                                                   boolean mustBeReference,
                                                                                   @NotNull PsiSwitchBlock context) {
      List<ExpectedTypeInfo> result = new ArrayList<>();
      Set<PsiType> processedTypes = new HashSet<>();
      for (PsiType expectedType : expectedTypes) {
        if (expectedType == null) {
          continue;
        }
        if (expectedType instanceof PsiClassType classType) {
          PsiClass resolved = classType.resolve();
          if (resolved != null && resolved.isEnum()) {
            processedTypes.add(expectedType);
            result.add(createInfoImpl(expectedType, expectedType));
            continue;
          }
        }
        if (expectedType.equalsToText(CommonClassNames.JAVA_LANG_STRING) ||
            TypeConversionUtil.isPrimitiveAndNotNullOrWrapper(expectedType)) {
          if (expectedType instanceof PsiPrimitiveType primitiveType) {
            if (primitiveType.equals(PsiTypes.longType()) ||
                primitiveType.equals(PsiTypes.doubleType()) ||
                primitiveType.equals(PsiTypes.floatType())) {
              return List.of(); //unexpected types, let's suggest default
            }
            if (mustBeReference) {
              expectedType = primitiveType.getBoxedType(context);
              if (expectedType == null) {
                continue;
              }
            } else {
              addWithWrapper(processedTypes, PsiTypes.intType(), result, context);
            }
          }
          addWithWrapper(processedTypes, expectedType, result, context);
          continue;
        }
        return List.of(); //something unexpected, let's suggest default
      }
      return result;
    }

    private static void addWithWrapper(@NotNull Set<PsiType> processedTypes,
                                       @NotNull PsiType expectedType,
                                       @NotNull List<ExpectedTypeInfo> result,
                                       @Nullable PsiElement context) {
      addIfNotExist(processedTypes, expectedType, result, createInfoImpl(expectedType, expectedType));
      if (context!=null && expectedType instanceof PsiPrimitiveType primitiveType) {
        PsiClassType type = primitiveType.getBoxedType(context);
        if (type != null) {
          addIfNotExist(processedTypes, type, result, createInfoImpl(type, type));
        }
      }
    }

