package com.jetbrains.python.documentation;

import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiReferenceBase;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.ArrayUtil;
import com.jetbrains.python.psi.PyDocStringOwner;
import com.jetbrains.python.psi.PyFunction;
import com.jetbrains.python.psi.PyNamedParameter;
import com.jetbrains.python.psi.impl.ParamHelper;
import org.jetbrains.annotations.NotNull;

import java.util.List;

/**
 * @author yole
 */
public class DocStringParameterReference extends PsiReferenceBase<PsiElement> {
  public DocStringParameterReference(PsiElement element, TextRange range) {
    super(element, range);
  }

  @Override
  public PsiElement resolve() {
    PyDocStringOwner owner = PsiTreeUtil.getParentOfType(getElement(), PyDocStringOwner.class);
    if (owner instanceof PyFunction) {
      final String paramName = getCanonicalText();
      return ((PyFunction) owner).getParameterList().getElementNamed(paramName);
    }
    return null;
  }

  @NotNull
  @Override
  public Object[] getVariants() {
    PyDocStringOwner owner = PsiTreeUtil.getParentOfType(getElement(), PyDocStringOwner.class);
    if (owner instanceof PyFunction) {
      final List<PyNamedParameter> namedParameters = ParamHelper.collectNamedParameters(((PyFunction)owner).getParameterList());
      return ArrayUtil.toObjectArray(namedParameters);
    }
    return ArrayUtil.EMPTY_OBJECT_ARRAY;
  }
}

@@ -234,24 +234,102 @@
 .ArrayUtil;%0A
+import com.jetbrains.python.PyNames;%0Aimport com.jetbrains.python.psi.PyClass;%0A
 import com.j
@@ -551,16 +551,59 @@
 NotNull;
+%0Aimport org.jetbrains.annotations.Nullable;
 %0A%0Aimport
@@ -1021,67 +1021,330 @@
     
-final String paramName = getCanonicalText();%0A      return (
+return resolveParameter((PyFunction)owner);%0A    %7D%0A    if (owner instanceof PyClass) %7B%0A      final PyFunction init = ((PyClass)owner).findMethodByName(PyNames.INIT, false);%0A      if (init != null) %7B%0A        return resolveParameter(init);%0A      %7D%0A    %7D%0A    return null;%0A  %7D%0A%0A  @Nullable%0A  private PsiElement resolveParameter
 (PyF
@@ -1354,16 +1354,34 @@
 tion
-)
  owner)
+ %7B%0A    return owner
 .get
@@ -1416,41 +1416,27 @@
 med(
-paramName);%0A    %7D%0A    return null
+getCanonicalText())
 ;%0A  

resolve @param tags in class doc comment to constructor parameters