package de.axp.hierarchyhighlighter;

import com.google.common.collect.Sets;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.editor.RangeMarker;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiMethod;
import de.axp.hierarchyhighlighter.MethodBackgroundPainter.PaintType;

import java.util.HashSet;
import java.util.Set;

class MethodHighlighter {

	private final MethodFinder methodFinder;
	private final MethodFolder methodFolder;
	private final MethodBackgroundPainter methodBackgroundPainter;

	MethodHighlighter() {
		methodFinder = new MethodFinder();
		methodFolder = new MethodFolder();
		methodBackgroundPainter = new MethodBackgroundPainter();
	}

	void highlight(Editor editor, PsiElement psiElement) {
		methodBackgroundPainter.getPaintJobs().forEach(RangeMarker::dispose);
		methodBackgroundPainter.getPaintJobs().clear();

		if (isValid(psiElement)) {
			PsiMethod psiMethod = (PsiMethod) psiElement;
			Set<PsiMethod> parentMethods = methodFinder.findMethodCallsOf(psiMethod);
			Set<PsiMethod> childMethods = methodFinder.findAllChildMethodsOf(psiMethod);

			foldMethodsInFile(editor, psiMethod);
			unfoldImportantMethods(editor, psiMethod, parentMethods, childMethods);
			paintBackgroundOfImportantMethods(editor, psiMethod, parentMethods, childMethods);
		}
	}

	private boolean isValid(PsiElement psiElement) {
		return psiElement != null && psiElement instanceof PsiMethod;
	}

	private void foldMethodsInFile(Editor editor, PsiMethod psiMethod) {
		methodFolder.foldMethods(editor, methodFinder.findAllChildMethodsOf(psiMethod.getContainingClass()));
	}

	private void unfoldImportantMethods(Editor editor, PsiMethod psiMethod, Set<PsiMethod> parentMethods, Set<PsiMethod> childMethods) {
		HashSet<PsiMethod> psiMethodsToUnfold = new HashSet<>();
		psiMethodsToUnfold.add(psiMethod);
		psiMethodsToUnfold.addAll(parentMethods);
		psiMethodsToUnfold.addAll(childMethods);
		methodFolder.unfoldMethod(editor, psiMethodsToUnfold);
	}

	private void paintBackgroundOfImportantMethods(Editor editor, PsiMethod psiMethod, Set<PsiMethod> parentMethods, Set<PsiMethod> childMethods) {
		methodBackgroundPainter.paintBackgroundOf(editor, PaintType.CURRENT_METHOD, Sets.newHashSet(psiMethod));
		methodBackgroundPainter.paintBackgroundOf(editor, PaintType.PARENT_METHOD, parentMethods);
		methodBackgroundPainter.paintBackgroundOf(editor, PaintType.CHILD_METHOD, childMethods);
	}
}

@@ -190,22 +190,147 @@
 lij.
-psi.PsiElement
+openapi.project.Project;%0D%0Aimport com.intellij.psi.PsiDocumentManager;%0D%0Aimport com.intellij.psi.PsiElement;%0D%0Aimport com.intellij.psi.PsiFile
 ;%0D%0Ai
@@ -460,24 +460,51 @@
 l.HashSet;%0D%0A
+import java.util.Objects;%0D%0A
 import java.
@@ -1072,16 +1072,52 @@
 Element)
+ && isInSameFile(editor, psiElement)
 ) %7B%0D%0A%09%09%09
@@ -1591,16 +1591,16 @@
 ent) %7B%0D%0A
-
 %09%09return
@@ -1654,32 +1654,363 @@
 iMethod;%0D%0A%09%7D%0D%0A%0D%0A
+%09private boolean isInSameFile(Editor editor, PsiElement psiElement) %7B%0D%0A%09%09Project project = editor.getProject();%0D%0A%09%09if (project != null) %7B%0D%0A%09%09%09PsiFile psiFile = PsiDocumentManager.getInstance(project).getPsiFile(editor.getDocument());%0D%0A%09%09%09return Objects.equals(psiFile, psiElement.getContainingFile());%0D%0A%09%09%7D%0D%0A%09%09return false;%0D%0A%09%7D%0D%0A%0D%0A
 %09private void fo

Assert that the handled method is in same PsiFile