package com.intellij.psi.util;

import com.intellij.psi.PsiElement;
import com.intellij.psi.search.PsiElementProcessor;

import java.util.ArrayList;

public class PsiTreeUtil {
  public static boolean isAncestor(PsiElement ancestor, PsiElement element, boolean strict) {
    PsiElement parent = strict ? element.getParent() : element;
    while (true) {
      if (parent == null) return false;
      if (parent.equals(ancestor)) return true;
      parent = parent.getParent();
    }
  }

  public static PsiElement findCommonParent(PsiElement element1, PsiElement element2) {
    ArrayList parents1 = getParents(element1);
    ArrayList parents2 = getParents(element2);
    int size = Math.min(parents1.size(), parents2.size());
    PsiElement parent = null;
    for (int i = 1; i <= size; i++) {
      PsiElement parent1 = (PsiElement) parents1.get(parents1.size() - i);
      PsiElement parent2 = (PsiElement) parents2.get(parents2.size() - i);
      if (!parent1.equals(parent2)) break;
      parent = parent1;
    }
    return parent;
  }

  private static ArrayList getParents(PsiElement element) {
    ArrayList parents = new ArrayList();
    PsiElement parent = element;
    while (parent != null) {
      parents.add(parent);
      parent = parent.getParent();
    }
    return parents;
  }

  public static PsiElement getChildOfType(PsiElement element, Class aClass) {
    for(PsiElement child = element.getFirstChild(); child != null; child = child.getNextSibling()){
      if (aClass.isInstance(child)) return child;
    }
    return null;
  }

  public static PsiElement getNextSiblingOfType(PsiElement sibling, Class aClass) {
    for(PsiElement child = sibling.getNextSibling(); child != null; child = child.getNextSibling()){
      if (aClass.isInstance(child)) return child;
    }
    return null;
  }

  public static PsiElement getPrevSiblingOfType(PsiElement sibling, Class aClass) {
    for(PsiElement child = sibling.getPrevSibling(); child != null; child = child.getPrevSibling()){
      if (aClass.isInstance(child)) return child;
    }
    return null;
  }

  public static PsiElement getParentOfType(PsiElement element, Class aClass) {
    return getParentOfType(element, aClass, true);
  }

  public static PsiElement getParentOfType(PsiElement element, Class aClass, boolean strict) {
    if (element == null) return null;
    if (strict) {
      element = element.getParent();
    }

    while (element != null && !aClass.isInstance(element)) {
      element = element.getParent();
    }

    return element;
  }

  public static PsiElement getParentOfType(PsiElement element, Class[] classes) {
    return getParentOfType(element, classes, true);
  }

  public static PsiElement getParentOfType(PsiElement element, Class[] classes, boolean strict) {
    if (strict) {
      element = element.getParent();
    }

    while (element != null) {
      for (int i = 0; i < classes.length; i++) {
        Class aClass = classes[i];
        if (aClass.isInstance(element)) return element;
      }
      element = element.getParent();
    }

    return element;
  }

  public static PsiElement[] collectElements(PsiElement element, PsiElementFilter filter) {
    PsiElementProcessor.CollectFilteredElements processor = new PsiElementProcessor.CollectFilteredElements(filter);
    processElements(element, processor);
    return processor.toArray();
  }

  public static boolean processElements(PsiElement element, PsiElementProcessor processor) {
    if (element == null) return true;
    if (!processor.execute(element)) return false;
    for (PsiElement child = element.getFirstChild(); child != null; child = child.getNextSibling()) {
      if (!processElements(child, processor)) return false;
    }

    return true;
  }
}