package com.jetbrains.python.codeInsight;

import com.intellij.psi.PsiElement;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyQualifiedName;
import com.jetbrains.python.psi.resolve.ResolveImportUtil;
import com.jetbrains.python.psi.stubs.PyClassNameIndex;
import org.jetbrains.annotations.Nullable;

import java.util.List;

/**
 * @author yole
 */
public abstract class PyPsiPath {
  @Nullable
  public abstract PsiElement resolve(PsiElement module);

  public static class ToFile extends PyPsiPath {
    private final PyQualifiedName myQualifiedName;

    public ToFile(String qualifiedName) {
      myQualifiedName = PyQualifiedName.fromDottedString(qualifiedName);
    }

    @Nullable
    @Override
    public PsiElement resolve(PsiElement context) {
      final List<PsiElement> elements = ResolveImportUtil.resolveModulesInRoots(myQualifiedName, context);
      return elements.size() > 0 ? elements.get(0) : null;
    }
  }

  public static class ToClassQName extends PyPsiPath {
    private final PyQualifiedName myQualifiedName;

    public ToClassQName(String qualifiedName) {
      myQualifiedName = PyQualifiedName.fromDottedString(qualifiedName);
    }

    @Nullable
    @Override
    public PsiElement resolve(PsiElement context) {
      return PyClassNameIndex.findClass(myQualifiedName.toString(), context.getProject());
    }
  }

  public static class ToClass extends PyPsiPath {
    private final PyPsiPath myParent;
    private final String myClassName;

    public ToClass(PyPsiPath parent, String className) {
      myParent = parent;
      myClassName = className;
    }

    @Override
    public PsiElement resolve(PsiElement context) {
      PsiElement parent = myParent.resolve(context);
      if (parent == null) {
        return null;
      }
      if (parent instanceof PyFile) {
        return ((PyFile) parent).findTopLevelClass(myClassName);
      }
      for (PsiElement element : parent.getChildren()) {
        if (element instanceof PyClass && myClassName.equals(((PyClass)element).getName())) {
          return element;
        }
      }
      return parent;
    }
  }

  public static class ToFunction extends PyPsiPath {
    private final PyPsiPath myParent;
    private final String myFunctionName;

    public ToFunction(PyPsiPath parent, String functionName) {
      myParent = parent;
      myFunctionName = functionName;
    }

    @Override
    public PsiElement resolve(PsiElement context) {
      PsiElement parent = myParent.resolve(context);
      if (parent == null) {
        return null;
      }
      if (parent instanceof PyFile) {
        return ((PyFile) parent).findTopLevelFunction(myFunctionName);
      }
      if (parent instanceof PyClass) {
        return ((PyClass) parent).findMethodByName(myFunctionName, false);
      }
      for (PsiElement element : parent.getChildren()) {
        if (element instanceof PyFunction && myFunctionName.equals(((PyFunction)element).getName())) {
          return element;
        }
      }
      return parent;
    }
  }

  public static class ToClassAttribute extends PyPsiPath {
    private final PyPsiPath myParent;
    private final String myAttributeName;

    public ToClassAttribute(PyPsiPath parent, String attributeName) {
      myAttributeName = attributeName;
      myParent = parent;
    }

    @Override
    public PsiElement resolve(PsiElement context) {
      PsiElement parent = myParent.resolve(context);
      if (!(parent instanceof PyClass)) {
        return null;
      }
      return ((PyClass)parent).findClassAttribute(myAttributeName, true);
    }
  }

  public static class ToCall extends PyPsiPath {
    private final PyPsiPath myParent;
    private final String myCallName;
    private final String[] myArgs;

    public ToCall(PyPsiPath parent, String callName, String... args) {
      myParent = parent;
      myCallName = callName;
      myArgs = args;
    }

    @Override
    public PsiElement resolve(PsiElement context) {
      PsiElement parent = myParent.resolve(context);
      if (parent == null) {
        return null;
      }
      CallFinder finder = new CallFinder(myCallName, myArgs);
      parent.accept(finder);
      return finder.myResult != null ? finder.myResult : parent;
    }
  }

  private static class CallFinder extends PyRecursiveElementVisitor {
    private PsiElement myResult;
    private final String myCallName;
    private final String[] myArgs;

    public CallFinder(String callName, String[] args) {
      myCallName = callName;
      myArgs = args;
    }

    @Override
    public void visitPyCallExpression(PyCallExpression node) {
      if (myResult != null) {
        return;
      }
      super.visitPyCallExpression(node);

      final PyExpression callee = node.getCallee();
      if (callee instanceof PyReferenceExpression) {
        String calleeName = ((PyReferenceExpression) callee).getReferencedName();
        if (myCallName.equals(calleeName)) {
          final PyExpression[] args = node.getArguments();
          if (myArgs.length <= args.length) {
            boolean argsMatch = true;
            for (int i = 0; i < myArgs.length; i++) {
              if (!(args[i] instanceof PyStringLiteralExpression) ||
                  !myArgs [i].equals(((PyStringLiteralExpression)args[i]).getStringValue())) {
                argsMatch = false;
                break;
              }
            }
            if (argsMatch) {
              myResult = node;
            }
          }
        }
      }
    }
  }
}

@@ -1895,32 +1895,73 @@
 sName);%0A      %7D%0A
+      if (parent instanceof PyClass) %7B%0A  
       for (PsiEl
@@ -1992,32 +1992,34 @@
 etChildren()) %7B%0A
+  
         if (elem
@@ -2098,32 +2098,34 @@
 ))) %7B%0A          
+  
 return element;%0A
@@ -2124,34 +2124,38 @@
 lement;%0A        
-%7D%0A
+  %7D%0A  
       %7D%0A      re
@@ -2148,38 +2148,563 @@
     %7D%0A      
-return parent;
+%7D%0A      ClassFinder finder = new ClassFinder(myClassName);%0A      parent.acceptChildren(finder);%0A      return finder.myResult != null ? finder.myResult : parent;%0A    %7D%0A  %7D%0A%0A  private static class ClassFinder extends PyRecursiveElementVisitor %7B%0A    private final String myName;%0A    private PyClass myResult;%0A%0A    public ClassFinder(String name) %7B%0A      myName = name;%0A    %7D%0A%0A    @Override%0A    public void visitPyClass(PyClass node) %7B%0A      super.visitPyClass(node);%0A      if (myName.equals(node.getName())) %7B%0A        myResult = node;%0A      %7D
 %0A    %7D%0A  %7D%0A%0A

support resolving to classes inside a function