// Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

package com.intellij.plugins.thrift.editor;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Optional;

import com.google.common.base.Stopwatch;
import com.intellij.codeInsight.daemon.RelatedItemLineMarkerInfo;
import com.intellij.codeInsight.daemon.RelatedItemLineMarkerProvider;
import com.intellij.codeInsight.navigation.NavigationGutterIconBuilder;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.plugins.thrift.ThriftBundle;
import com.intellij.plugins.thrift.index.ThriftDeclarationIndex;
import com.intellij.plugins.thrift.lang.psi.ThriftDeclaration;
import com.intellij.psi.PsiClass;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiIdentifier;
import com.intellij.psi.PsiModifierList;
import com.intellij.psi.search.GlobalSearchScope;

import org.jetbrains.annotations.NotNull;

import icons.ThriftIcons;

public class GoToThriftDefinitionMarkerProvider extends RelatedItemLineMarkerProvider {
  static Logger logger = Logger.getInstance(GoToThriftDefinition.class);

  @Override
  protected void collectNavigationMarkers(@NotNull PsiElement element,
                                          @NotNull Collection<? super RelatedItemLineMarkerInfo<?>> result) {
    if (element instanceof PsiClass) {
      PsiClass psiClass = (PsiClass) element;
      String name = psiClass.getName();
      if (name != null) {
        Stopwatch stopWatch = Stopwatch.createStarted();
        List<ThriftDeclaration> thriftDefinitions =
            ThriftDeclarationIndex.findDeclaration(name, psiClass.getProject(),
                GlobalSearchScope.allScope(psiClass.getProject()));
        Duration time = stopWatch.elapsed();
        if (time.toMillis() > 20) {
          logger.warn("GoToThriftDefinition action update took " + time.toMillis() + "ms");
        }
        if (!thriftDefinitions.isEmpty()) {
          NavigationGutterIconBuilder<PsiElement> builder =
              NavigationGutterIconBuilder.create(ThriftIcons.STRUCT)
                  .setTargets(thriftDefinitions)
                  .setTooltipText(ThriftBundle.message("thrift.goto.source"));
          Optional<PsiElement> identifier =
              Arrays.stream(element.getChildren()).filter(x -> x instanceof PsiIdentifier).findFirst();
          // Some languages (Scala) don't put PsiIdentifier here, PsiModifierList is a fallback
          Optional<PsiElement> modifiersList =
              Arrays.stream(element.getChildren()).filter(x -> x instanceof PsiModifierList).findFirst();
          if (identifier.isPresent()) {
            result.add(builder.createLineMarkerInfo(identifier.get()));
          } else if (modifiersList.isPresent()) {
            result.add(builder.createLineMarkerInfo(modifiersList.get()));
          }
        }
      }
    }
  }
}

@@ -302,16 +302,38 @@
 ptional;
+%0Aimport java.util.Set;
 %0A%0Aimport
@@ -367,16 +367,55 @@
 pwatch;%0A
+import com.google.common.collect.Sets;%0A
 import c
@@ -1271,16 +1271,379 @@
 lass);%0A%0A
+  private static Set%3CString%3E thriftStructBaseClasses =%0A      Sets.newHashSet(%0A          %22com.twitter.scrooge.ThriftStruct%22,%0A          %22org.apache.thrift.TBase%22%0A      );%0A%0A  private boolean isThriftStruct(PsiClass element) %7B%0A    return Arrays.stream(element.getSupers()).anyMatch(superClass -%3E thriftStructBaseClasses.contains(superClass.getQualifiedName()));%0A  %7D%0A%0A
   @Overr
@@ -1862,16 +1862,53 @@
 PsiClass
+ && isThriftStruct((PsiClass)element)
 ) %7B%0A    

Optimize line marker provider (#93)