package com.theoryinpractice.testng.ui.actions;

import com.intellij.execution.CantRunException;
import com.intellij.execution.Executor;
import com.intellij.execution.Location;
import com.intellij.execution.actions.JavaRerunFailedTestsAction;
import com.intellij.execution.configurations.RunProfileState;
import com.intellij.execution.runners.ExecutionEnvironment;
import com.intellij.execution.testframework.AbstractTestProxy;
import com.intellij.execution.testframework.SourceScope;
import com.intellij.execution.testframework.TestConsoleProperties;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.module.Module;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.ui.ComponentContainer;
import com.intellij.openapi.util.Computable;
import com.intellij.psi.PsiClass;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiMethod;
import com.intellij.psi.PsiModifier;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.util.containers.ContainerUtil;
import com.theoryinpractice.testng.configuration.SearchingForTestsTask;
import com.theoryinpractice.testng.configuration.TestNGConfiguration;
import com.theoryinpractice.testng.configuration.TestNGRunnableState;
import com.theoryinpractice.testng.util.TestNGUtil;
import org.jetbrains.annotations.NotNull;

import java.io.File;
import java.net.ServerSocket;
import java.util.*;

public class RerunFailedTestsAction extends JavaRerunFailedTestsAction {
  public RerunFailedTestsAction(@NotNull ComponentContainer componentContainer, @NotNull TestConsoleProperties consoleProperties) {
    super(componentContainer, consoleProperties);
  }

  @Override
  protected MyRunProfile getRunProfile(@NotNull ExecutionEnvironment environment) {
    final TestNGConfiguration configuration = (TestNGConfiguration)getModel().getProperties().getConfiguration();
    final List<AbstractTestProxy> failedTests = getFailedTests(configuration.getProject());
    return new MyRunProfile(configuration) {
      @Override
      @NotNull
      public Module[] getModules() {
        return configuration.getModules();
      }

      @Override
      public RunProfileState getState(@NotNull Executor executor, @NotNull ExecutionEnvironment env) {
        return new TestNGRunnableState(env, configuration) {
          @Override
          protected SearchingForTestsTask createSearchingForTestsTask(ServerSocket serverSocket,
                                                                      final TestNGConfiguration config, final File tempFile) {
            return new SearchingForTestsTask(serverSocket, config, tempFile, client) {
              @Override
              protected void fillTestObjects(final Map<PsiClass, Collection<PsiMethod>> classes) throws CantRunException {
                final HashMap<PsiClass, Collection<PsiMethod>> fullClassList = ContainerUtil.newHashMap();
                super.fillTestObjects(fullClassList);
                for (final PsiClass aClass : fullClassList.keySet()) {
                  if (!ApplicationManager.getApplication().runReadAction(new Computable<Boolean>() {
                    @Override
                    public Boolean compute() {
                      return TestNGUtil.hasTest(aClass);
                    }
                  })) {
                    classes.put(aClass, fullClassList.get(aClass));
                  }
                }

                final GlobalSearchScope scope = config.getConfigurationModule().getSearchScope();
                final Project project = config.getProject();
                for (AbstractTestProxy proxy : failedTests) {
                  final Location location = proxy.getLocation(project, scope);
                  if (location != null) {
                    final PsiElement element = location.getPsiElement();
                    if (element instanceof PsiMethod && element.isValid()) {
                      final PsiMethod psiMethod = (PsiMethod)element;
                      PsiClass psiClass = psiMethod.getContainingClass();
                      if (psiClass != null && psiClass.hasModifierProperty(PsiModifier.ABSTRACT)) {
                        final AbstractTestProxy parent = proxy.getParent();
                        final PsiElement elt = parent != null ? parent.getLocation(project, scope).getPsiElement() : null;
                        if (elt instanceof PsiClass) {
                          psiClass = (PsiClass)elt;
                        }
                      }
                      Collection<PsiMethod> psiMethods = classes.get(psiClass);
                      if (psiMethods == null) {
                        psiMethods = new ArrayList<PsiMethod>();
                        classes.put(psiClass, psiMethods);
                      }
                      psiMethods.add(psiMethod);
                    }
                  }
                }
              }
            };
          }
        };
      }
    };
  }
}

@@ -3596,16 +3596,22 @@
    for (
+final 
 Abstract
@@ -3641,24 +3641,674 @@
 dTests) %7B%0A  
+                ApplicationManager.getApplication().runReadAction(new Runnable() %7B%0A                    public void run() %7B%0A                      includeFailedTestWithDependencies(classes, scope, project, proxy);%0A                    %7D%0A                  %7D);%0A                %7D%0A              %7D%0A%0A              private void includeFailedTestWithDependencies(Map%3CPsiClass, Collection%3CPsiMethod%3E%3E classes,%0A                                                             GlobalSearchScope scope,%0A                                                             Project project,%0A                                                             AbstractTestProxy proxy) %7B%0A
             
@@ -4384,26 +4384,24 @@
             
-  
 if (location
@@ -4422,34 +4422,32 @@
                 
- 
- 
 final PsiElement
@@ -4479,26 +4479,24 @@
 iElement();%0A
-  
             
@@ -4570,34 +4570,32 @@
                 
-  
 final PsiMethod 
@@ -4646,18 +4646,16 @@
         
- 
- 
 PsiClass
@@ -4694,26 +4694,24 @@
 ingClass();%0A
-  
             
@@ -4792,26 +4792,24 @@
 BSTRACT)) %7B%0A
-  
             
@@ -4884,34 +4884,32 @@
                 
-  
 final PsiElement
@@ -5009,26 +5009,24 @@
             
-  
 if (elt inst
@@ -5040,26 +5040,24 @@
 PsiClass) %7B%0A
-  
             
@@ -5108,38 +5108,34 @@
                 
-  %7D%0A  
+%7D%0A
                 
@@ -5160,18 +5160,16 @@
         
- 
- 
 Collecti
@@ -5210,34 +5210,32 @@
 .get(psiClass);%0A
-  
                 
@@ -5278,34 +5278,32 @@
                 
-  
 psiMethods = new
@@ -5341,34 +5341,32 @@
                 
- 
- 
 classes.put(psiC
@@ -5396,38 +5396,34 @@
                 
-  %7D%0A  
+%7D%0A
                 
@@ -5453,38 +5453,16 @@
 ethod);%0A
-                    %7D%0A
         

read action during search in background