package com.theoryinpractice.testng.ui.actions;

import com.intellij.execution.CantRunException;
import com.intellij.execution.ExecutionException;
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.openapi.module.Module;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.ui.ComponentContainer;
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.theoryinpractice.testng.configuration.SearchingForTestsTask;
import com.theoryinpractice.testng.configuration.TestNGConfiguration;
import com.theoryinpractice.testng.configuration.TestNGRunnableState;
import org.jetbrains.annotations.NotNull;

import java.io.File;
import java.net.ServerSocket;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

public class RerunFailedTestsAction extends JavaRerunFailedTestsAction {

  public RerunFailedTestsAction(@NotNull ComponentContainer componentContainer) {
    super(componentContainer);
  }

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

      public RunProfileState getState(@NotNull Executor executor, @NotNull ExecutionEnvironment env) throws ExecutionException {

        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 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);
                    }
                  }
                }
              }
            };
          }
        };
      }
    };
  }

}

@@ -471,16 +471,76 @@
 tProxy;%0A
+import com.intellij.openapi.application.ApplicationManager;%0A
 import c
@@ -570,24 +570,24 @@
 ule.Module;%0A
-
 import com.i
@@ -666,24 +666,69 @@
 tContainer;%0A
+import com.intellij.openapi.util.Computable;%0A
 import com.i
@@ -1115,24 +1115,76 @@
 nableState;%0A
+import com.theoryinpractice.testng.util.TestNGUtil;%0A
 import org.j
@@ -1286,91 +1286,9 @@
 til.
-ArrayList;%0Aimport java.util.Collection;%0Aimport java.util.List;%0Aimport java.util.Map
+*
 ;%0A%0Ap
@@ -2415,32 +2415,32 @@
       @Override%0A
-
               pr
@@ -2546,16 +2546,625 @@
 ption %7B%0A
+                final HashMap%3CPsiClass, Collection%3CPsiMethod%3E%3E fullClassList = new HashMap%3C%3E();%0A                super.fillTestObjects(fullClassList);%0A                for (final PsiClass aClass : fullClassList.keySet()) %7B%0A                  if (!ApplicationManager.getApplication().runReadAction(new Computable%3CBoolean%3E() %7B%0A                    @Override%0A                    public Boolean compute() %7B%0A                      return TestNGUtil.hasTest(aClass);%0A                    %7D%0A                  %7D)) %7B%0A                    classes.put(aClass, fullClassList.get(aClass));%0A                  %7D%0A                %7D%0A%0A
         

include all config-only classes in rerun failed (IDEA-124571)