package com.jetbrains.python.testing;

import com.intellij.execution.Location;
import com.intellij.execution.PsiLocation;
import com.intellij.openapi.progress.ProgressManager;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.io.FileUtil;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.testIntegration.TestLocationProvider;
import com.jetbrains.python.psi.PyClass;
import com.jetbrains.python.psi.PyFunction;
import com.jetbrains.python.psi.stubs.PyClassNameIndex;
import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class PythonUnitTestTestIdUrlProvider implements TestLocationProvider {
  @NonNls
  private static final String PROTOCOL_ID = "python_uttestid";

  @NotNull
  public List<Location> getLocation(@NotNull final String protocolId, @NotNull final String path,
                                    final Project project) {
    if (!PROTOCOL_ID.equals(protocolId)) {
      return Collections.emptyList();
    }

    final List<String> list = StringUtil.split(path, ".");
    if (list.isEmpty()) {
      return Collections.emptyList();
    }
    final int listSize = list.size();

    // parse path as [ns.]*fileName.className.methodName
    if (listSize < 3) {
      return Collections.emptyList();
    }

    final String className = list.get(listSize - 2);
    final String methodName = list.get(listSize - 1);

    String fileName = list.get(listSize - 3);
    if (fileName.indexOf("%") >= 0) {
      fileName = fileName.substring(0, fileName.lastIndexOf("%"));
    }

    final List<Location> locations = new ArrayList<Location>();
    for (PyClass cls : PyClassNameIndex.find(className, project, false)) {
      ProgressManager.checkCanceled();

      final PyFunction method = cls.findMethodByName(methodName, true);
      if (method == null) {
        continue;
      }

      final String clsFileName = FileUtil.getNameWithoutExtension(cls.getContainingFile().getName());
      if (!clsFileName.equalsIgnoreCase(fileName)) {
        continue;
      }

      locations.add(new PsiLocation<PyFunction>(project, method));
    }

    return locations;
  }
}

@@ -591,16 +591,59 @@
 NotNull;
+%0Aimport org.jetbrains.annotations.Nullable;
 %0A%0Aimport
@@ -1342,16 +1342,17 @@
 lassName
+%5B
 .methodN
@@ -1354,16 +1354,18 @@
 thodName
+%5D%0A
 %0A    if 
@@ -1378,11 +1378,12 @@
 ize 
-%3C 3
+== 2
 ) %7B%0A
@@ -1399,40 +1399,96 @@
 urn 
-Collections.emptyList();%0A    %7D%0A%0A
+findLocations(project, list.get(0), list.get(1), null);%0A    %7D%0A    if (listSize %3E 2) %7B%0A  
     
@@ -1536,16 +1536,18 @@
 e - 2);%0A
+  
     fina
@@ -1597,16 +1597,18 @@
 );%0A%0A    
+  
 String f
@@ -1641,16 +1641,582 @@
 e - 3);%0A
+      final List%3CLocation%3E locations = findLocations(project, fileName, className, methodName);%0A      if (locations.size() %3E 0) %7B%0A        return locations;%0A      %7D%0A      return findLocations(project, list.get(listSize-2), list.get(listSize-1), null);%0A    %7D%0A    return Collections.emptyList();%0A  %7D%0A%0A%0A  private static List%3CLocation%3E findLocations(Project project,%0A                                              String fileName,%0A                                              String className,%0A                                              @Nullable String methodName) %7B%0A
     if (
@@ -2502,135 +2502,8 @@
 );%0A%0A
-      final PyFunction method = cls.findMethodByName(methodName, true);%0A      if (method == null) %7B%0A        continue;%0A      %7D%0A%0A
     
@@ -2675,25 +2675,277 @@
 ue;%0A      %7D%0A
-%0A
+      if (methodName == null) %7B%0A        locations.add(new PsiLocation%3CPyClass%3E(project, cls));%0A      %7D%0A      else %7B%0A        final PyFunction method = cls.findMethodByName(methodName, true);%0A        if (method == null) %7B%0A          continue;%0A        %7D%0A%0A  
       locati
@@ -2995,24 +2995,32 @@
 , method));%0A
+      %7D%0A
     %7D%0A%0A    r

test location can be class as well as method