// Copyright 2000-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.intellij.util;

import com.intellij.openapi.application.Application;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.util.ThrowableComputable;
import com.intellij.openapi.util.registry.Registry;
import com.intellij.util.containers.ContainerUtil;
import org.jetbrains.annotations.NotNull;

import java.util.HashSet;
import java.util.Set;

public final class SlowOperations {

  private static final Logger LOG = Logger.getInstance(SlowOperations.class);
  private static final Set<String> ourReportedTraces = new HashSet<>();
  private static final String[] misbehavingFrames = {
    "org.jetbrains.kotlin.idea.refactoring.introduce.introduceVariable.KotlinIntroduceVariableHandler",
    "org.jetbrains.kotlin.idea.actions.KotlinAddImportAction",
    "org.jetbrains.kotlin.idea.codeInsight.KotlinCopyPasteReferenceProcessor",
    "com.intellij.apiwatcher.plugin.presentation.bytecode.UsageHighlighter",
  };
  private static boolean ourAllowedFlag = System.getenv("TEAMCITY_VERSION") != null ||
                                          !ApplicationManager.getApplication().isInternal();

  private SlowOperations() {}

  /**
   * If you get an exception from this method, then you need to move the computation to the background
   * while also trying to avoid blocking the UI thread as well. It's okay if the API changes in the process,
   * e.g. instead of wrapping implementation of some extension into {@link #allowSlowOperations},
   * it's better to admit that the EP semantic as a whole requires index access,
   * and then move the iteration over all extensions to the background on the platform-side.
   * <p/>
   * In cases when it's impossible to do so, the computation should be wrapped into {@code allowSlowOperations} explicitly.
   * This way it's possible to find all such operations by searching usages of {@code allowSlowOperations}.
   *
   * @see com.intellij.openapi.application.NonBlockingReadAction
   * @see com.intellij.openapi.application.CoroutinesKt#readAction
   * @see com.intellij.openapi.actionSystem.ex.ActionUtil#underModalProgress
   */
  public static void assertSlowOperationsAreAllowed() {
    if (ourAllowedFlag) {
      return;
    }
    if (Registry.is("ide.enable.slow.operations.in.edt")) {
      return;
    }
    Application application = ApplicationManager.getApplication();
    if (application.isUnitTestMode() || !application.isDispatchThread() || application.isWriteAccessAllowed()) {
      return;
    }
    String stackTrace = ExceptionUtil.currentStackTrace();
    if (ContainerUtil.or(misbehavingFrames, stackTrace::contains) || !ourReportedTraces.add(stackTrace)) {
      return;
    }
    LOG.error("Slow operations are prohibited in the EDT");
  }

  public static <T, E extends Throwable> T allowSlowOperations(@NotNull ThrowableComputable<T, E> computable) throws E {
    if (ourAllowedFlag || !ApplicationManager.getApplication().isDispatchThread()) {
      return computable.compute();
    }
    ourAllowedFlag = true;
    try {
      return computable.compute();
    }
    finally {
      ourAllowedFlag = false;
    }
  }

  public static <E extends Throwable> void allowSlowOperations(@NotNull ThrowableRunnable<E> runnable) throws E {
    allowSlowOperations(() -> {
      runnable.run();
      return null;
    });
  }
}

@@ -158,24 +158,77 @@
 llij.util;%0A%0A
+import com.intellij.openapi.application.AccessToken;%0A
 import com.i
@@ -532,16 +532,90 @@
 erUtil;%0A
+import com.intellij.util.ui.EDT;%0Aimport org.jetbrains.annotations.NonNls;%0A
 import o
@@ -1284,16 +1284,22 @@
  static 
+final 
 boolean 
@@ -1303,25 +1303,25 @@
 an ourAl
-lowedFlag
+waysAllow
  = Syste
@@ -1360,16 +1360,126 @@
 null %7C%7C%0A
+                                                ApplicationManager.getApplication().isUnitTestMode() %7C%7C%0A      
         
@@ -1562,16 +1562,57 @@
 ernal();
+%0A  private static long ourAllowanceCount;
 %0A%0A  priv
@@ -2657,25 +2657,25 @@
 f (ourAl
-lowedFlag
+waysAllow
 ) %7B%0A    
@@ -2849,40 +2849,8 @@
 if (
-application.isUnitTestMode() %7C%7C 
 !app
@@ -2879,16 +2879,57 @@
 ead() %7C%7C
+%0A        ourAllowanceCount %3E 0 %7C%7C%0A       
  applica
@@ -3359,250 +3359,101 @@
     
-if (ourAllowedFlag %7C%7C !ApplicationManager.getApplication().isDispatchThread()) %7B%0A      return computable.compute();%0A    %7D%0A    ourAllowedFlag = true;%0A    try %7B%0A      return computable.compute();%0A    %7D%0A    finally %7B%0A      ourAllowedFlag = false
+try (AccessToken ignore = allowSlowOperations(%22generic%22)) %7B%0A      return computable.compute()
 ;%0A  
@@ -3578,16 +3578,42 @@
  E %7B%0A   
+ try (AccessToken ignore =
  allowSl
@@ -3629,13 +3629,19 @@
 ons(
-() -%3E
+%22generic%22))
  %7B%0A 
@@ -3669,29 +3669,457 @@
     
-  return null;%0A    %7D)
+%7D%0A  %7D%0A%0A  /** @noinspection unused*/%0A  @NotNull%0A  public static AccessToken allowSlowOperations(@NotNull @NonNls String activityName) %7B%0A    if (ourAlwaysAllow %7C%7C !EDT.isCurrentThreadEdt()) %7B%0A      return AccessToken.EMPTY_ACCESS_TOKEN;%0A    %7D%0A    ourAllowanceCount++;%0A    return new AccessToken() %7B%0A      @Override%0A      public void finish() %7B%0A        //noinspection AssignmentToStaticFieldFromInstanceMethod%0A        ourAllowanceCount--;%0A      %7D%0A    %7D
 ;%0A  

introduce `AccessToken`-based API to reduce stack size