package com.zzsn.knowbase.util.tree;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.ObjectUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author lkg
 * @description: 构造树形结构工具类
 * @date 2021/12/15 14:18
 */
public class TreeUtil {

    /**
     * 单个树
     *
     * @param list 数据集合
     * @param node 根节点数据
     * @author lkg
     * @date 2021/12/17 16:33
     */
    public static <T extends Node> void tree(List<T> list, T node) {
        String id = node.getId();
        List<T> tree = tree(list, id);
        node.setChildren(tree);
    }

    /**
     * 树 集合
     *
     * @param list   数据集合
     * @param rootId 根节点id
     * @author lkg
     * @date 2021/12/17 16:33
     */
    public static <T extends Node> List<T> tree(List<T> list, String rootId) {
        Map<String, List<T>> treeMap = new HashMap<>();
        for (T item : list) {
            String parentId = item.getPid();
            List<T> children;
            if (treeMap.containsKey(parentId)) {
                children = treeMap.get(parentId);
            } else {
                children = new ArrayList<>();
                treeMap.put(parentId, children);
            }
            children.add(item);
        }
        return getTreeByMap(treeMap, rootId);
    }

    /**
     * 根据当前节点获取所有上级节点
     *
     * @param list        数据集合
     * @param currentNode 当前节点
     * @param flag        是否包含自己本身(true-是；false-否)
     * @author lkg
     * @date 2022/12/27 11:32
     */
    public static <T extends Node> LinkedHashSet<T> getUpList(List<T> list, T currentNode, boolean flag) {
        if (ObjectUtils.isNotEmpty(currentNode)) {
            LinkedHashSet<T> treeSet = new LinkedHashSet<>();
            if (flag) {
                treeSet.add(currentNode);
            }
            String pid = currentNode.getPid();
            List<T> parentNodes = list.stream().filter(node -> node.getId().equals(pid)).collect(Collectors.toList());
            if (CollectionUtils.isNotEmpty(parentNodes)) {
                T parentNode = parentNodes.get(0);
                treeSet.add(parentNode);
                LinkedHashSet<T> upTreeSet = getUpList(list, parentNode, flag);
                if (CollectionUtils.isNotEmpty(upTreeSet)) {
                    treeSet.addAll(upTreeSet);
                }
            }
            return treeSet;
        }
        return null;
    }

    /**
     * 获取某节点下所有节点id得集合
     *
     * @param list:   数据集合
     * @param rootId: 节点id
     * @param flag:   是否包含自己本身(true-是；false-否)
     * @author lkg
     * @date 2023/2/23
     */
    public static <T extends Node> List<String> belowList(List<T> list, String rootId, boolean flag) {
        List<String> nodes = new ArrayList<>();
        if (flag) {
            nodes.add(rootId);
        }
        Map<String, List<T>> treeMap = new HashMap<>();
        for (T item : list) {
            String parentId = item.getPid();
            List<T> children;
            if (treeMap.containsKey(parentId)) {
                children = treeMap.get(parentId);
            } else {
                children = new ArrayList<>();
                treeMap.put(parentId, children);
            }
            children.add(item);
        }
        getBelowList(treeMap, rootId, nodes);
        return nodes;
    }

    private static <T extends Node> List<T> getBelowList(Map<String, List<T>> treeMap, String parentId, List<String> belowList) {
        List<T> children = treeMap.get(parentId);
        if (children == null) {
            return null;
        }
        for (T item : children) {
            belowList.add(item.getId());
            item.setChildren(getBelowList(treeMap, item.getId(), belowList));
        }
        return children;
    }


    private static <T extends Node> List<T> getTreeByMap(Map<String, List<T>> treeMap, String parentId) {
        List<T> children = treeMap.get(parentId);
        if (children == null) {
            return null;
        }
        for (Node item : children) {
            item.setChildren(getTreeByMap(treeMap, item.getId()));
        }
        return children;
    }

}
