package com.kaylerrenslow.armaplugin.lang;

import com.intellij.lang.ASTNode;
import com.intellij.openapi.fileTypes.FileType;
import com.intellij.openapi.project.Project;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiFileFactory;
import com.intellij.psi.TokenType;
import com.intellij.psi.tree.IElementType;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.function.Function;

/**
 * @author Kayler
 * @since 09/06/2017
 */
public class PsiUtil {

	/**
	 * Creates and returns a PsiFile with the given text of the given file type.
	 *
	 * @throws ClassCastException when the PsiFile created couldn't be cast to T
	 */
	@NotNull
	@SuppressWarnings("unchecked")
	public static <T extends PsiFile> T createFile(@NotNull Project project, @NotNull String text, @NotNull FileType fileType) {
		String fileName = "fake_sqf_file.sqf";
		return (T) PsiFileFactory.getInstance(project).createFileFromText(fileName, fileType, text);
	}

	/**
	 * Creates a new Psi file via {@link #createFile(Project, String, FileType)} and then gets the first element in the entire
	 * file where the PsiElement's class is an instance of the given class
	 *
	 * @param project project
	 * @param text    file text content
	 * @param ft      the FileType
	 * @param clazz   first PsiElement class to find
	 * @return the element, or null if couldn't be found
	 */
	@Nullable
	public static <T extends PsiElement> T createElement(@NotNull Project project, @NotNull String text, @NotNull FileType ft, @NotNull Class<T> clazz) {
		PsiFile f = createFile(project, text, ft);
		return findFirstDescendantElement(f, clazz);
	}

	@NotNull
	public static ASTNode getFirstDescendantNode(@NotNull PsiElement element) {
		ASTNode cursor = element.getNode();
		while (cursor.getFirstChildNode() != null) {
			cursor = cursor.getFirstChildNode();
		}
		return cursor;
	}

	/**
	 * Will traverse up the AST tree and send each discovered node into the callback, including the starting node.
	 * The function can return true, which will then terminate the traversal, or return false or null to continue the traversal.
	 *
	 * @param start    where to start
	 * @param callback function invoked for each discovered node
	 */
	public static void traverseUp(@NotNull ASTNode start, @NotNull Function<ASTNode, Boolean> callback) {
		Boolean stop = callback.apply(start);
		if (stop != null && stop) {
			return;
		}
		ASTNode parent = start;
		while (parent != null) {
			parent = parent.getTreeParent();
			stop = callback.apply(parent);
			if (stop != null && stop) {
				return;
			}
		}
	}

	/**
	 * Traverses the entire ast tree with BFS, starting from start. Each node that is found will be sent to callback.
	 * It is also possible to stop the traversal at any time with callback by returning true in it
	 *
	 * @param start    starting ASTNode
	 * @param callback function that returns true to null/false to end search completely, or false to continue search
	 */
	public static void traverseBreadthFirstSearch(@NotNull ASTNode start, @NotNull Function<ASTNode, Boolean> callback) {
		Boolean stop = callback.apply(start);
		if (stop != null && stop) {
			return;
		}
		ASTNode[] children = start.getChildren(null);
		LinkedList<ASTNode> nodes = new LinkedList<>();

		for (ASTNode child : children) {
			nodes.addLast(child);
		}
		ASTNode node;
		while (nodes.size() > 0) {
			node = nodes.removeFirst();
			stop = callback.apply(node);
			if (stop != null && stop) {
				return;
			}
			children = node.getChildren(null);
			for (ASTNode child : children) {
				nodes.addLast(child);
			}
		}
	}

	/**
	 * Traverses the entire ast tree with DFS, starting from start. Each node that is found will be sent to callback.
	 * It is also possible to stop the traversal at any time with callback by returning true in it
	 *
	 * @param start    starting ASTNode
	 * @param callback function that returns true to end search completely, or null/false to continue search
	 */
	public static void traverseDepthFirstSearch(@NotNull ASTNode start, @NotNull Function<ASTNode, Boolean> callback) {
		Boolean stop = callback.apply(start);
		if (stop != null && stop) {
			return;
		}
		ASTNode[] children = start.getChildren(null);
		LinkedList<ASTNode> nodes = new LinkedList<>();

		for (ASTNode child : children) {
			nodes.push(child);
		}
		ASTNode node;
		while (nodes.size() > 0) {
			node = nodes.pop();
			stop = callback.apply(node);
			if (stop != null && stop) {
				return;
			}
			children = node.getChildren(null);
			for (ASTNode child : children) {
				nodes.push(child);
			}
		}
	}

	/**
	 * Traverses the entire ast tree with BFS, starting from start. Each node that is found will be sent to callback.
	 * It is also possible to stop the traversal at any time with callback by returning true in it. When returning true,
	 * the traversal will prevent the current node's children from being traversed.
	 * <p>
	 * This method will not terminate until all discovered nodes are traversed. If the callback returns true, there may
	 * still be nodes left to traverse because of the current node's parent had multiple children.
	 *
	 * @param start    starting ASTNode
	 * @param callback TraversalObjectFinder
	 */
	public static void traverseInLayers(@NotNull ASTNode start, @NotNull Function<ASTNode, Boolean> callback) {
		Boolean stop = callback.apply(start);
		if (stop != null && stop) {
			return;
		}
		ASTNode[] children = start.getChildren(null);
		LinkedList<ASTNode> nodes = new LinkedList<>();

		for (ASTNode child : children) {
			nodes.addLast(child);
		}
		ASTNode node;
		while (nodes.size() > 0) {
			node = nodes.removeFirst();
			stop = callback.apply(node);
			if (stop != null && stop) {
				continue;
			}
			children = node.getChildren(null);
			for (ASTNode child : children) {
				nodes.addLast(child);
			}
		}
	}

	/**
	 * Gets the closest next sibling, that is non-whitespace, relative to node
	 *
	 * @param node node to find sibling of
	 * @return non-whitespace sibling, or null if none was found
	 */
	@Nullable
	public static ASTNode getNextSiblingNotWhitespace(@NotNull ASTNode node) {
		return getNextSiblingNotType(node, TokenType.WHITE_SPACE);
	}

	/**
	 * Gets the closest next sibling, where the type is not skip, relative to node
	 *
	 * @param node node to find sibling of
	 * @param skip the token to skip
	 * @return non-skip sibling, or null if none was found
	 */
	@Nullable
	public static ASTNode getNextSiblingNotType(@NotNull ASTNode node, @NotNull IElementType skip) {
		ASTNode sibling = node.getTreeNext();
		while (sibling != null) {
			if (sibling.getElementType() == skip) {
				sibling = sibling.getTreeNext();
			} else {
				break;
			}
		}
		return sibling;
	}

	/**
	 * Gets the closest previous sibling, that is non-whitespace, relative to node
	 *
	 * @param node node to find sibling of
	 * @return non-whitespace sibling, or null if none was found
	 */
	@Nullable
	public static ASTNode getPrevSiblingNotWhitespace(@NotNull ASTNode node) {
		return getPrevSiblingNotType(node, TokenType.WHITE_SPACE);
	}

	/**
	 * Gets the closest previous sibling, that is not skip, relative to node
	 *
	 * @param node node to find sibling of
	 * @param skip what element type to skip
	 * @return non-whitespace sibling, or null if none was found
	 */
	@Nullable
	public static ASTNode getPrevSiblingNotType(@NotNull ASTNode node, @NotNull IElementType skip) {
		ASTNode sibling = node.getTreePrev();
		while (sibling != null) {
			if (sibling.getElementType() == skip) {
				sibling = sibling.getTreePrev();
			} else {
				break;
			}
		}
		return sibling;
	}

	/**
	 * Checks if the given node is a descendant of the given IElementType.<br>
	 * If textContent is not null, this method will also check if the ancestor is of correct type and ancestor's text is equal to textContent.
	 *
	 * @param node        node to check if has a ancestor of IElementType type
	 * @param type        IElementType to check
	 * @param textContent null if to disregard text of ancestor, otherwise check if ancestor's text is equal to textContent
	 * @return true if node has ancestor of IElementType type and ancestor's text matches textContent. If textContent is null, text can be anything for ancestor.
	 */
	public static boolean isDescendantOf(@NotNull ASTNode node, @NotNull IElementType type, @Nullable String textContent) {
		return getFirstAncestorOfType(node, type, textContent) != null;
	}

	/**
	 * Checks if the given node has an ancestor of the given IElementType. If there is one, this method will return that ancestor. Otherwise, it will return null.<br>
	 * If textContent is not null, this method will also check if the ancestor is of correct type and ancestor's text is equal to textContent.
	 *
	 * @param node        node to check if has a parent of IElementType type
	 * @param type        IElementType to check
	 * @param textContent null if to disregard text of ancestor, otherwise check if ancestor's text is equal to textContent
	 * @return node's ancestor if ancestor is of IElementType type if node's ancestor's text matches textContent. If textContent is null, text can be anything for ancestor.
	 */
	@Nullable
	public static ASTNode getFirstAncestorOfType(@NotNull ASTNode node, @NotNull IElementType type, @Nullable String textContent) {
		ASTNode parent = node.getTreeParent();
		boolean isChild = false;
		while (parent != null && !isChild) {
			parent = parent.getTreeParent();
			if (parent == null) {
				break;
			}
			isChild = parent.getElementType() == type && (textContent == null || parent.getText().equals(textContent));
		}
		return parent;
	}


	/**
	 * Checks if the given ASTNode is of IElementType et
	 *
	 * @param node ASTNode (if null, returns false)
	 * @param et   IElement type
	 * @return true if node is of type et, false otherwise
	 */
	public static boolean isOfElementType(@Nullable ASTNode node, @NotNull IElementType et) {
		return node != null && node.getElementType() == et;
	}

	/**
	 * Checks if the given PsiElement is of IElementType et
	 *
	 * @param pe PsiElement (if null, returns false)
	 * @param et IElement type
	 * @return true if pe is of type et, false otherwise
	 */
	public static boolean isOfElementType(@Nullable PsiElement pe, @NotNull IElementType et) {
		return pe != null && isOfElementType(pe.getNode(), et);
	}

	@Nullable
	public static <T extends PsiElement> T findFirstDescendantElement(@NotNull PsiElement element, @NotNull Class<T> type) {
		PsiElement child = element.getFirstChild();
		while (child != null) {
			if (type.isInstance(child)) {
				return (T) child;
			}
			T e = findFirstDescendantElement(child, type);
			if (e != null) {
				return e;
			}
			child = child.getNextSibling();
		}
		return null;
	}


	@NotNull
	public static <E extends PsiElement> List<E> findDescendantElementsOfInstance(@NotNull PsiElement rootElement,
																				  @NotNull Class<E> type,
																				  @Nullable PsiElement cursor,
																				  @Nullable String textContent) {
		ArrayList<E> list = new ArrayList<>();
		findDescendantElementsOfInstance(rootElement, type, cursor, textContent, list);
		return list;
	}

	private static <E extends PsiElement> void findDescendantElementsOfInstance(@NotNull PsiElement rootElement,
																				@NotNull Class<?> type,
																				@Nullable PsiElement cursor,
																				@Nullable String textContent,
																				@NotNull List<E> list) {
		PsiElement child = rootElement.getFirstChild();
		while (child != null) {
			if (cursor != null && child == cursor) {
				continue;
			}
			if (type.isAssignableFrom(child.getClass()) && (textContent == null || child.getText().equals(textContent))) {
				list.add((E) child);
			}
			findDescendantElementsOfInstance(child, type, cursor, textContent, list);
			child = child.getNextSibling();
		}
	}


	/**
	 * Get children of the given PsiElement that extend/are type of the given class
	 *
	 * @param element  element to get children of
	 * @param psiClass class
	 * @return list of all children
	 */
	@NotNull
	public static <T extends PsiElement> List<T> findChildrenOfType(@NotNull PsiElement element, @NotNull Class<T> psiClass) {
		List<T> list = new ArrayList<T>();
		PsiElement[] children = element.getChildren();
		for (PsiElement child : children) {
			if (psiClass.isInstance(child)) {
				list.add((T) child);
			}
		}
		return list;
	}
}
