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 org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Collection;
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 : getClassesByName(project, className)) {
      ProgressManager.checkCanceled();

      final PyFunction method = locateMethodInHierarchy(cls, methodName);
      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;
  }

  @Nullable
  private static PyFunction locateMethodInHierarchy(final PyClass cls, final String methodName) {
    PyFunction func = cls.findMethodByName(methodName, false);
    if (func != null) {
      return func;
    }

    for (PyClass ancestors : cls.iterateAncestors()) {
      func = ancestors.findMethodByName(methodName, false);
      if (func != null) {
        return func;
      }
    }

    return null;
  }

  private static Collection<PyClass> getClassesByName(final Project project, final String name) {
    return PyClassNameIndex.find(name, project, false); 
  }
}

@@ -596,107 +596,35 @@
 ll;%0A
-import org.jetbrains.annotations.Nullable;%0A%0Aimport java.util.ArrayList;%0Aimport java.util.Collection
+%0Aimport java.util.ArrayList
 ;%0Aim
@@ -1733,33 +1733,49 @@
 s cls : 
-get
+Py
 Class
-esBy
+NameIndex.find(class
 Name
-(
+, 
 project,
@@ -1775,24 +1775,20 @@
 roject, 
-classNam
+fals
 e)) %7B%0A  
@@ -1861,37 +1861,29 @@
 d = 
-locateMethodInHierarchy(cls, 
+cls.findMethodByName(
 meth
@@ -1888,16 +1888,22 @@
 thodName
+, true
 );%0A     
@@ -2237,589 +2237,7 @@
 %0A  %7D
-%0A%0A  @Nullable%0A  private static PyFunction locateMethodInHierarchy(final PyClass cls, final String methodName) %7B%0A    PyFunction func = cls.findMethodByName(methodName, false);%0A    if (func != null) %7B%0A      return func;%0A    %7D%0A%0A    for (PyClass ancestors : cls.iterateAncestors()) %7B%0A      func = ancestors.findMethodByName(methodName, false);%0A      if (func != null) %7B%0A        return func;%0A      %7D%0A    %7D%0A%0A    return null;%0A  %7D%0A%0A  private static Collection%3CPyClass%3E getClassesByName(final Project project, final String name) %7B%0A    return PyClassNameIndex.find(name, project, false); %0A  %7D
 %0A%7D%0A

simplify code