package org.asciidoc.intellij.findUsages;

import static com.intellij.patterns.PlatformPatterns.psiElement;
import static com.intellij.patterns.PlatformPatterns.psiFile;
import com.intellij.openapi.util.TextRange;
import com.intellij.patterns.PsiElementPattern;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiReference;
import com.intellij.psi.PsiReferenceContributor;
import com.intellij.psi.PsiReferenceProvider;
import com.intellij.psi.PsiReferenceRegistrar;
import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReferenceSet;
import com.intellij.util.ProcessingContext;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.asciidoc.intellij.lexer.AsciiDocTokenTypes;
import org.asciidoc.intellij.psi.AsciiDocFile;
import org.asciidoc.intellij.psi.AsciiDocLink;
import org.asciidoc.intellij.psi.AsciiDocRef;
import org.asciidoc.intellij.psi.AsciiDocReference;
import org.jetbrains.annotations.NotNull;

public class AsciiDocReferenceContributor extends PsiReferenceContributor {
  @Override
  public void registerReferenceProviders(@NotNull PsiReferenceRegistrar registrar) {

    final PsiElementPattern.Capture<AsciiDocRef> refCapture =
      psiElement(AsciiDocRef.class).inFile(psiFile(AsciiDocFile.class));

    registrar.registerReferenceProvider(refCapture,
      new PsiReferenceProvider() {
        @NotNull
        @Override
        public PsiReference[] getReferencesByElement(@NotNull PsiElement element,
                                                     @NotNull ProcessingContext
                                                       context) {
          int start = 0;
          PsiElement child = element.getFirstChild();
          while (child != null && child.getNode().getElementType() != AsciiDocTokenTypes.REF) {
            start += child.getTextLength();
            child = child.getNextSibling();
          }
          if (child != null) {
            return new PsiReference[]{
              new AsciiDocReference(element, TextRange.create(start, start + child.getTextLength()))
            };
          } else {
            return new PsiReference[]{};
          }
        }
      });

    final PsiElementPattern.Capture<AsciiDocLink> linkCapture =
      psiElement(AsciiDocLink.class).inFile(psiFile(AsciiDocFile.class));

    registrar.registerReferenceProvider(linkCapture,
      new PsiReferenceProvider() {
        @NotNull
        @Override
        public PsiReference[] getReferencesByElement(@NotNull PsiElement element,
          @NotNull ProcessingContext
            context) {
          List<PsiReference> references = new ArrayList<>();

          PsiReference anchorReference = findAnchor(element);
          if(anchorReference!=null) {
            references.add(anchorReference);
          }

          List<PsiReference> fileReferences = findFileReferences(element);

          references.addAll(fileReferences);

          return references.toArray(new PsiReference[0]);
        }
      });
  }

  private PsiReference findAnchor(PsiElement element) {
    int start = 0;
    PsiElement child = element.getFirstChild();
    while (child != null && child.getNode().getElementType() != AsciiDocTokenTypes.LINKANCHOR) {
      start += child.getTextLength();
      child = child.getNextSibling();
    }
    return child!=null ? new AsciiDocReference(element, TextRange.create(start, start + child.getTextLength())) : null;
  }

  private List<PsiReference> findFileReferences(PsiElement element) {
    int start = 0;
    PsiElement child = element.getFirstChild();
    while (child != null && child.getNode().getElementType() != AsciiDocTokenTypes.LINKFILE) {
      start += child.getTextLength();
      child = child.getNextSibling();
    }
    if (child != null) {
      return Arrays.asList(new FileReferenceSet(child.getText(), element, start, null, false).getAllReferences());
    } else {
      return Collections.emptyList();
    }
  }
}

@@ -40,135 +40,8 @@
 s;%0A%0A
-import static com.intellij.patterns.PlatformPatterns.psiElement;%0Aimport static com.intellij.patterns.PlatformPatterns.psiFile;%0A
 impo
@@ -80,16 +80,16 @@
 tRange;%0A
+
 import c
@@ -344,16 +344,100 @@
 istrar;%0A
+import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReference;%0A
 import c
@@ -563,114 +563,8 @@
 xt;%0A
-import java.util.ArrayList;%0Aimport java.util.Arrays;%0Aimport java.util.Collections;%0Aimport java.util.List;%0A
 impo
@@ -802,24 +802,24 @@
 cReference;%0A
-
 import org.j
@@ -847,16 +847,294 @@
 NotNull;
+%0Aimport org.jetbrains.annotations.Nullable;%0A%0Aimport java.util.ArrayList;%0Aimport java.util.Arrays;%0Aimport java.util.Collections;%0Aimport java.util.List;%0A%0Aimport static com.intellij.patterns.PlatformPatterns.psiElement;%0Aimport static com.intellij.patterns.PlatformPatterns.psiFile;
 %0A%0Apublic
@@ -3593,24 +3593,937 @@
  null;%0A  %7D%0A%0A
+  class LinkFileReferenceSet extends FileReferenceSet %7B%0A    LinkFileReferenceSet(String str, @NotNull PsiElement element, int startInElement, @Nullable PsiReferenceProvider provider, boolean isCaseSensitive) %7B%0A      super(str, element, startInElement, provider, isCaseSensitive);%0A    %7D%0A%0A    @Override%0A    protected void reparse() %7B%0A      super.reparse();%0A      if (myReferences.length %3E 0) %7B%0A        FileReference fileReference = myReferences%5BmyReferences.length - 1%5D;%0A        // if the reference can't find a file, try with an added '.adoc' extension%0A        if (fileReference.resolve() == null) %7B%0A          FileReference withExtension = createFileReference(fileReference.getRangeInElement(), fileReference.getIndex(), fileReference.getText() + %22.adoc%22);%0A          if (withExtension.resolve() != null) %7B%0A            myReferences%5BmyReferences.length - 1%5D = withExtension;%0A          %7D%0A        %7D%0A      %7D%0A    %7D%0A  %7D%0A%0A
   private Li
@@ -4880,16 +4880,21 @@
 ist(new 
+ Link
 FileRefe

allow file references without .adoc to be found