/*
 * Copyright (c) 2018-present, Jim Kynde Meyer
 * All rights reserved.
 * <p>
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
package com.intellij.lang.jsgraphql.ide.project.javascript;

import com.intellij.lang.injection.InjectedLanguageManager;
import com.intellij.lang.jsgraphql.ide.injection.javascript.GraphQLLanguageInjectionUtil;
import com.intellij.lang.jsgraphql.ide.project.GraphQLInjectionSearchHelper;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiManager;
import com.intellij.psi.PsiRecursiveElementVisitor;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.util.indexing.FileBasedIndex;

import java.util.Collections;
import java.util.function.Consumer;

public class GraphQLJavascriptInjectionSearchHelper implements GraphQLInjectionSearchHelper {

    @Override
    public boolean isJSGraphQLLanguageInjectionTarget(PsiElement host) {
        return GraphQLLanguageInjectionUtil.isJSGraphQLLanguageInjectionTarget(host);
    }

    /**
     * Uses the {@link GraphQLInjectionIndex} to process injected GraphQL PsiFiles
     *
     * @param scopedElement the starting point of the enumeration settings the scopedElement of the processing
     * @param schemaScope   the search scope to use for limiting the schema definitions
     * @param consumer      a consumer that will be invoked for each injected GraphQL PsiFile
     */
    public void processInjectedGraphQLPsiFiles(PsiElement scopedElement, GlobalSearchScope schemaScope, Consumer<PsiFile> consumer) {
        final PsiManager psiManager = PsiManager.getInstance(scopedElement.getProject());
        final InjectedLanguageManager injectedLanguageManager = InjectedLanguageManager.getInstance(scopedElement.getProject());
        FileBasedIndex.getInstance().getFilesWithKey(GraphQLInjectionIndex.NAME, Collections.singleton(GraphQLInjectionIndex.DATA_KEY), virtualFile -> {
            final PsiFile fileWithInjection = psiManager.findFile(virtualFile);
            if (fileWithInjection != null) {
                fileWithInjection.accept(new PsiRecursiveElementVisitor() {
                    @Override
                    public void visitElement(PsiElement element) {
                        if (GraphQLLanguageInjectionUtil.isJSGraphQLLanguageInjectionTarget(element)) {
                            injectedLanguageManager.enumerate(element, (injectedPsi, places) -> {
                                consumer.accept(injectedPsi);
                            });
                        } else {
                            // visit deeper until injection found
                            super.visitElement(element);
                        }
                    }
                });
            }
            return true;
        }, schemaScope);

    }
}

@@ -488,24 +488,84 @@
 archHelper;%0A
+import com.intellij.openapi.project.IndexNotReadyException;%0A
 import com.i
@@ -1681,24 +1681,42 @@
 consumer) %7B%0A
+        try %7B%0A    
         fina
@@ -1797,24 +1797,28 @@
 ));%0A        
+   
+ 
 final Inject
@@ -1918,32 +1918,36 @@
 .getProject());%0A
+    
         FileBase
@@ -2087,32 +2087,36 @@
 %3E %7B%0A            
+ 
+   
 final PsiFile fi
@@ -2167,16 +2167,20 @@
 lFile);%0A
+    
         
@@ -2232,16 +2232,20 @@
         
+    
 fileWith
@@ -2312,24 +2312,28 @@
             
+   
+ 
 @Override%0A  
@@ -2326,24 +2326,28 @@
   @Override%0A
+    
             
@@ -2397,24 +2397,28 @@
  element) %7B%0A
+    
             
@@ -2505,24 +2505,28 @@
 element)) %7B%0A
+    
             
@@ -2643,16 +2643,20 @@
         
+    
 consumer
@@ -2697,36 +2697,44 @@
                 
-%7D);%0A
+    %7D);%0A    
                 
@@ -2778,16 +2778,20 @@
         
+  
+  
 // visit
@@ -2820,16 +2820,20 @@
 n found%0A
+    
         
@@ -2897,34 +2897,41 @@
                 
-%7D%0A
+    %7D%0A   
                 
@@ -2926,34 +2926,39 @@
                 
+ 
 %7D%0A
+    
                 
@@ -2965,34 +2965,42 @@
 %7D);%0A            
-%7D%0A
+    %7D%0A    
             retu
@@ -3016,16 +3016,20 @@
         
+    
 %7D, schem
@@ -3037,16 +3037,132 @@
 Scope);%0A
+        %7D catch (IndexNotReadyException e) %7B%0A            // can't search yet (e.g. during project startup)%0A        %7D
 %0A    %7D%0A%7D

Handle indexing exception, e.g. during project startup (#164)