// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.intellij.sh.codeInsight;

import com.intellij.openapi.util.TextRange;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiReferenceBase;
import com.intellij.sh.psi.ShFunctionDefinition;
import com.intellij.sh.psi.ShFunctionName;
import com.intellij.sh.psi.ShGenericCommandDirective;
import com.intellij.util.containers.ContainerUtil;
import com.intellij.util.containers.hash.LinkedHashMap;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.List;
import java.util.Map;

public class ShFunctionReference extends PsiReferenceBase<PsiElement> {
  private final PsiFile myFile;
  private Map<String, FunctionContext> myExecutionContext;

  public ShFunctionReference(@NotNull PsiElement element) {
    super(element, TextRange.create(0, element.getTextLength()));
    myFile = myElement.getContainingFile();
  }

  @Nullable
  @Override
  public PsiElement resolve() {
    if (myFile == null) return null;
    myExecutionContext = new LinkedHashMap<>();

    // Execution context should be build, not check inner function if it was not called
    PsiElement target = checkChildren(myFile, false);
    if (target != null) return target;
    return checkUnvisitedMethods();
  }

  @Nullable
  private PsiElement checkUnvisitedMethods() {
    List<FunctionContext> notVisitedFunctions = ContainerUtil.filter(myExecutionContext.values(), it -> !it.visited);
    // In this case we should visit all not visited function one by one with checking inner function also
    for (FunctionContext functionContext : notVisitedFunctions) {
      functionContext.visited = true;
      PsiElement target = checkChildren(functionContext.function, true);
      if (target != null) return target;
    }
    return null;
  }

  private PsiElement checkChildren(PsiElement element, boolean visitInnerFunctions) {
    if (element == null) return null;
    for (PsiElement child : element.getChildren()) {
      PsiElement target = buildExecutionContext(skipUnnecessaryNodes(child), visitInnerFunctions);
      if (target != null) return target;
    }
    return null;
  }

  private PsiElement buildExecutionContext(PsiElement element, boolean visitInnerFunctions) {
    if (element == null) return null;
    if (element instanceof ShFunctionDefinition) {
      // If we meet function definition, should add it to the context like a not visited node
      return handleFunction((ShFunctionDefinition)element, visitInnerFunctions);
    }
    if (element instanceof ShGenericCommandDirective) {
      // Check if command linked to function which we searching for. If linked function is not visited, we should check
      // all its entire function and add them to the execution context
      if (myExecutionContext.containsKey(element.getText())) {
        FunctionContext functionContext = myExecutionContext.get(element.getText());
        if (visitInnerFunctions) {
          if (element == myElement && functionContext.visited) return functionContext.function.getFunctionName();
        } else {
          if (element == myElement) return functionContext.function.getFunctionName();
          if (!functionContext.visited) {
            functionContext.visited = true;
            PsiElement target = checkChildren(functionContext.function, visitInnerFunctions);
            if (target != null) return target;
          }
        }
      }
    }
    return checkChildren(element, visitInnerFunctions);
  }

  private static PsiElement skipUnnecessaryNodes(@NotNull PsiElement element) {
    PsiElement nextElement = element;
    while (shouldBeSkipped(nextElement)) {
      nextElement = nextElement.getChildren()[0];
    }
    return nextElement;
  }

  private static boolean shouldBeSkipped(@NotNull PsiElement element) {
    return element.getChildren().length == 1 && !(element instanceof ShGenericCommandDirective || element instanceof ShFunctionDefinition);
  }

  private PsiElement handleFunction(@NotNull ShFunctionDefinition functionDefinition, boolean visitInnerFunctions) {
    ShFunctionName functionName = functionDefinition.getFunctionName();
    assert functionName != null;
    FunctionContext functionContext = new FunctionContext(functionDefinition);
    myExecutionContext.put(functionName.getText(), functionContext);
    if (!visitInnerFunctions) return null;
    functionContext.visited = true;
    return checkChildren(functionDefinition, visitInnerFunctions);
  }

  private static class FunctionContext {
    private final ShFunctionDefinition function;
    private boolean visited;

    private FunctionContext(ShFunctionDefinition function) {
      this.function = function;
    }
  }
}

@@ -473,16 +473,54 @@
 ective;%0A
+import com.intellij.sh.psi.ShLiteral;%0A
 import c
@@ -2944,20 +2944,179 @@
 t%0A      
-if (
+ShGenericCommandDirective genericCommand = (ShGenericCommandDirective)element;%0A      ShLiteral literal = genericCommand.getLiteral();%0A      if (literal != null && 
 myExecut
@@ -3138,23 +3138,23 @@
 ainsKey(
-element
+literal
 .getText
@@ -3225,23 +3225,23 @@
 ext.get(
-element
+literal
 .getText
@@ -3286,39 +3286,39 @@
 %7B%0A          if (
-element
+literal
  == myElement &&
@@ -3421,31 +3421,31 @@
         if (
-element
+literal
  == myElemen

Fix function reference resolving