package org.coursera.courier;

import com.intellij.lang.cacheBuilder.DefaultWordsScanner;
import com.intellij.lang.cacheBuilder.WordsScanner;
import com.intellij.lang.findUsages.FindUsagesProvider;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiNamedElement;
import com.intellij.psi.tree.TokenSet;
import org.coursera.courier.psi.CourierTypeNameDeclaration;
import org.coursera.courier.psi.CourierTypes;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public class CourierFindUsagesProvider implements FindUsagesProvider {
  private static final DefaultWordsScanner WORDS_SCANNER =
    new DefaultWordsScanner(new CourierLexerAdapter(),
      TokenSet.create(CourierTypes.IDENTIFIER),
      TokenSet.create(
        CourierTypes.LINE_COMMENT,
        CourierTypes.BLOCK_COMMENT_NON_EMPTY,
        CourierTypes.BLOCK_COMMENT_EMPTY,
        CourierTypes.COMMA),
      TokenSet.create(
        CourierTypes.JSON_STRING,
        CourierTypes.JSON_NUMBER,
        CourierTypes.JSON_BOOLEAN,
        CourierTypes.JSON_NULL));

  @Nullable
  @Override
  public WordsScanner getWordsScanner() {
    return WORDS_SCANNER;
  }

  @Override
  public boolean canFindUsagesFor(@NotNull PsiElement psiElement) {
    return psiElement instanceof PsiNamedElement;
  }

  @Nullable
  @Override
  public String getHelpId(@NotNull PsiElement psiElement) {
    return null;
  }

  @NotNull
  @Override
  public String getType(@NotNull PsiElement element) {
    if (element instanceof CourierTypeNameDeclaration) {
      return "type"; // TODO: lookup the data schema type
    } else {
      return "";
    }
  }

  @NotNull
  @Override
  public String getDescriptiveName(@NotNull PsiElement element) {
    if (element instanceof CourierTypeNameDeclaration) {
      return ((CourierTypeNameDeclaration)element).getFullname().toString();
    } else {
      return "";
    }
  }

  @NotNull
  @Override
  public String getNodeText(@NotNull PsiElement element, boolean useFullName) {
    if (element instanceof CourierTypeNameDeclaration) {
      CourierTypeNameDeclaration declaration = (CourierTypeNameDeclaration) element;
      String name = declaration.getName();
      if (name != null) {
        return name;
      } else {
        return "";
      }
    } else {
      return "";
    }
  }
}

@@ -570,16 +570,17 @@
 vider %7B%0A
+%0A
   privat
@@ -598,101 +598,31 @@
 nal 
-DefaultWordsScanner WORDS_SCANNER =%0A    new DefaultWordsScanner(new CourierLexerAdapter(),%0A  
+TokenSet identifiers =%0A
     
@@ -661,23 +661,61 @@
 NTIFIER)
-,%0A     
+;%0A%0A  private static final TokenSet comments =
  TokenSe
@@ -716,36 +716,32 @@
 okenSet.create(%0A
-    
     CourierTypes
@@ -755,28 +755,24 @@
 OMMENT,%0A    
-    
 CourierTypes
@@ -793,36 +793,32 @@
 _NON_EMPTY,%0A    
-    
 CourierTypes.BLO
@@ -835,20 +835,16 @@
 _EMPTY,%0A
-    
     Cour
@@ -858,23 +858,61 @@
 s.COMMA)
-,%0A     
+;%0A%0A  private static final TokenSet literals =
  TokenSe
@@ -921,20 +921,16 @@
 create(%0A
-    
     Cour
@@ -951,20 +951,16 @@
 STRING,%0A
-    
     Cour
@@ -981,20 +981,16 @@
 NUMBER,%0A
-    
     Cour
@@ -1012,20 +1012,16 @@
 OOLEAN,%0A
-    
     Cour
@@ -1039,17 +1039,16 @@
 ON_NULL)
-)
 ;%0A%0A  @Nu
@@ -1116,28 +1116,154 @@
     
-return WORDS_SCANNER
+// See https://devnet.jetbrains.com/message/5537369%0A    return new DefaultWordsScanner(new CourierLexerAdapter(), identifiers, comments, literals)
 ;%0A  

Fix how IntelliJ plugin indexes code for find-usages.