package com.niuan.erp.common.utils;

import com.niuan.erp.common.base.BaseTree;

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

public class TreeUtils {

    public static <T> List<BaseTree> generateTree(List<T> nodeList,
                                            Function<T, Object> getIdFunc,
                                            Function<T, String> getLabelFunc,
                                            Function<T, Object> getParentIdFunc,
                                            Object rootId) {
        HashMap<Object, BaseTree> map = new HashMap<>();
        List<BaseTree> rootList = new ArrayList<>();
        nodeList.forEach(node -> {
            Object id = getIdFunc.apply(node);
            String label = getLabelFunc.apply(node);
            Object parentId = getParentIdFunc.apply(node);
            BaseTree treeNode = new BaseTree(id, label, null);
            // 检查是不是顶级 Node
            if (rootId.equals(parentId)) {
                rootList.add(treeNode);
            }
            // 添加父节点
            if (map.containsKey(parentId)) {
                if (map.get(parentId).getChildren() != null) map.get(parentId).getChildren().add(treeNode);
                else map.get(parentId).setChildren(new ArrayList<>(List.of(treeNode)));
            } else {
                map.put(parentId, new BaseTree(parentId, null, new ArrayList<>(List.of(treeNode))));
            }

            // 判断自己
            if (!map.containsKey(id)) {
                map.put(id, treeNode);
            } else {
                treeNode.setChildren(map.get(id).getChildren());
                map.put(id, treeNode);
            }
        });
        System.out.println(rootList);
        return rootList;
    }

}