package uk.co.reecedunn.intellij.plugin.xquery.resolve.reference;

import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiReferenceBase;
import org.jetbrains.annotations.NotNull;
import uk.co.reecedunn.intellij.plugin.xquery.ast.xquery.*;
import uk.co.reecedunn.intellij.plugin.xquery.functional.Option;
import uk.co.reecedunn.intellij.plugin.xquery.psi.PsiNavigation;
import uk.co.reecedunn.intellij.plugin.xquery.psi.XQueryPrologResolver;

public class XQueryFunctionNameReference extends PsiReferenceBase<XQueryEQName> {
    public XQueryFunctionNameReference(XQueryEQName element, TextRange range) {
        super(element, range);
    }

    @Override
    public PsiElement resolve() {
        return getElement().resolvePrefixNamespace().map((ns) -> {
            if (ns.getDeclaration() instanceof XQueryPrologResolver) {
                XQueryPrologResolver provider = (XQueryPrologResolver) ns.getDeclaration();
                Option<XQueryProlog> prolog = provider.resolveProlog();
                if (!prolog.isDefined()) {
                    return null;
                }

                PsiElement parent = getElement().getParent();
                int arity = 0;
                if (parent instanceof XQueryFunctionCall) {
                    arity = ((XQueryFunctionCall) parent).getArity();
                } else if (parent instanceof XQueryNamedFunctionRef) {
                    arity = ((XQueryNamedFunctionRef) parent).getArity();
                }

                CharSequence localName = getElement().getLocalNameElement().getText();
                PsiElement annotation = prolog.get().getFirstChild();
                while (annotation != null) {
                    if (annotation instanceof XQueryAnnotatedDecl) {
                        XQueryFunctionDecl functionDecl = PsiNavigation.findChildByClass(annotation, XQueryFunctionDecl.class);
                        if (functionDecl != null) {
                            XQueryEQName functionName = PsiNavigation.findChildByClass(functionDecl, XQueryEQName.class);
                            if (functionName != null && functionName.getLocalNameElement().getText().equals(localName)) {
                                if (functionDecl.getArity() == arity) {
                                    return functionName;
                                }
                            }
                        }
                    }

                    annotation = annotation.getNextSibling();
                }
            }
            return null;
        }).getOrElse(null);
    }

    @NotNull
    @Override
    public Object[] getVariants() {
        return new Object[0];
    }
}

@@ -1124,32 +1124,133 @@
              %7D%0A%0A
+                return getElement().getLocalName().map(PsiElement::getText).map((localName) -%3E %7B%0A    
                 
@@ -1311,16 +1311,20 @@
         
+    
 int arit
@@ -1330,16 +1330,20 @@
 ty = 0;%0A
+    
         
@@ -1406,32 +1406,36 @@
                 
+    
 arity = ((XQuery
@@ -1460,32 +1460,36 @@
 nt).getArity();%0A
+    
                 
@@ -1555,32 +1555,36 @@
                 
+    
 arity = ((XQuery
@@ -1633,39 +1633,39 @@
             
-%7D%0A%0A
     
+%7D%0A%0A
         
     CharSequ
@@ -1656,83 +1656,8 @@
     
-    CharSequence localName = getElement().getLocalNameElement().getText();%0A
     
@@ -1656,36 +1656,32 @@
                 
-    
 PsiElement annot
@@ -1714,24 +1714,28 @@
 rstChild();%0A
+    
             
@@ -1783,24 +1783,28 @@
             
+    
 if (annotati
@@ -1860,24 +1860,28 @@
             
+   
+ 
 XQueryFuncti
@@ -1988,32 +1988,36 @@
                 
+    
 if (functionDecl
@@ -2052,24 +2052,28 @@
             
+   
+ 
 XQueryEQName
@@ -2174,36 +2174,75 @@
                 
-if (
+    Option%3CPsiElement%3E functionLocalName = 
 functionName != 
@@ -2242,18 +2242,33 @@
 ame 
-!
+=
 = null 
-&&
+? Option.none() :
  fun
@@ -2293,49 +2293,153 @@
 Name
+();%0A                                if (functionLocalName.map(Psi
 Element
-().
+::
 getText
-().equals(localName)) %7B%0A
+).map((name) -%3E name.equals(localName)).getOrElse(false)) %7B%0A    
     
@@ -2538,24 +2538,28 @@
             
+    
 return funct
@@ -2567,16 +2567,54 @@
 onName;%0A
+                                    %7D%0A
         
@@ -2695,39 +2695,18 @@
       %7D%0A
-                    %7D%0A
 %0A
+ 
         
@@ -2717,16 +2717,19 @@
         
+   
 annotati
@@ -2770,33 +2770,106 @@
                 
-%7D
+    %7D%0A                    return null;%0A                %7D).getOrElse(null);
 %0A            %7D%0A 

Use the Option version of getLocalName.