import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 泛型生成树工具
 *
 * @param <T> 实体类型
 * @author lazydong
 */
public class GenericsTree<T> {
    private Builder<T> builder;

    public static <T> GenericsTree.Builder<T> createBuilder(Class<T> tClass) {
        return new GenericsTree.Builder<T>();
    }

    public List<T> buildTree() {
        int currentLevel = 1;
        // 获取根节点数据
        if (builder.rootList == null) {
            Comparator<T> currentSortCondition = getCurrentLevelSortCondition(currentLevel);
            if (currentSortCondition != null) {
                builder.rootList = builder.dataList.stream()
                        .filter(temp -> builder.rootLevelCondition.apply(temp))
                        .sorted(currentSortCondition)
                        .collect(Collectors.toList());
            } else {
                builder.rootList = builder.dataList.stream()
                        .filter(temp -> builder.rootLevelCondition.apply(temp))
                        .collect(Collectors.toList());
            }
        }
        if (!builder.rootList.isEmpty()) {
            // 父节点分组
            Map<Object, List<T>> groupingMap = builder.dataList.stream().collect(Collectors.groupingBy(builder.parentFiled));
            fillMiddleLevel(groupingMap, builder.rootList, currentLevel + 1);
        }
        return builder.rootList;
    }

    /**
     * 循环填充中间层级
     *
     * @param groupingMap       分组数据
     * @param currentParentList 当前父节点
     * @param currentLevel      当前层级
     */
    private void fillMiddleLevel(Map<Object, List<T>> groupingMap, List<T> currentParentList, int currentLevel) {
        for (T t : currentParentList) {
            Object apply = builder.childFiled.apply(t);
            List<T> ts = groupingMap.get(apply);
            if (ts == null) {
                continue;
            }
            if (builder.maxDepth == 0 || builder.maxDepth > currentLevel) {
                fillMiddleLevel(groupingMap, ts, currentLevel + 1);
                Comparator<T> currentLevelSortCondition = getCurrentLevelSortCondition(currentLevel);
                if (currentLevelSortCondition != null) {
                    ts = ts.stream().sorted(currentLevelSortCondition).collect(Collectors.toList());
                }
                builder.fillChildNodeMethod.accept(t, ts);
            } else {
                // 已经到了最后一级了
                fillLastLevel(groupingMap, currentParentList, currentLevel);
            }
        }
    }

    /**
     * 填充最后一层
     *
     * @param groupingMap       分组数据
     * @param currentParentList 当前父节点
     * @param currentLevel      当前层级
     */
    private void fillLastLevel(Map<Object, List<T>> groupingMap, List<T> currentParentList, int currentLevel) {
        Comparator<T> currentLevelSortCondition = getCurrentLevelSortCondition(currentLevel);
        for (T t : currentParentList) {
            Object apply = builder.childFiled.apply(t);
            List<T> ts = groupingMap.get(apply);
            if (ts == null) {
                continue;
            }
            List<T> resultList = new ArrayList<>(ts);
            if (builder.lastLevelMerge) {
                loopsFillLastLevel(groupingMap, ts, resultList);
            }
            if (currentLevelSortCondition != null) {
                resultList = resultList.stream().sorted(currentLevelSortCondition).collect(Collectors.toList());
            }
            builder.fillChildNodeMethod.accept(t, resultList);
        }
    }

    /**
     * 循环填充最后一层数据
     *
     * @param groupingMap       分组数据
     * @param currentParentList 当前父节点
     * @param resultList        最后一层数据
     */
    private void loopsFillLastLevel(Map<Object, List<T>> groupingMap, List<T> currentParentList, List<T> resultList) {
        for (T t : currentParentList) {
            Object apply = builder.childFiled.apply(t);
            List<T> ts = groupingMap.get(apply);
            if (ts == null) {
                continue;
            }
            resultList.addAll(ts);
            loopsFillLastLevel(groupingMap, ts, resultList);
        }
    }

    private Comparator<T> getCurrentLevelSortCondition(Integer currentLevel) {
        if (builder.sortLevelCondition.containsKey(currentLevel)) {
            return builder.sortLevelCondition.get(currentLevel);
        }
        return builder.sortCondition;
    }


    /**
     * 条件构建器
     *
     * @param <T>
     */
    static class Builder<T> {

        /**
         * 父节点字段
         */
        private Function<T, Object> parentFiled;

        /**
         * 子节点字段
         */
        private Function<T, Object> childFiled;

        /**
         * 外部传入的根节点
         */
        private List<T> rootList;

        /**
         * 生成根节点条件，rootList优先
         */
        private Function<T, Boolean> rootLevelCondition;

        /**
         * 统一排序条件，sortLevelCondition 优先
         */
        private Comparator<T> sortCondition;

        /**
         * 每一层级的排序条件
         */
        private final Map<Integer, Comparator<T>> sortLevelCondition = new HashMap<>();

        /**
         * 填充子节点的方法
         */
        private BiConsumer<T, List<T>> fillChildNodeMethod;

        /**
         * 树的最大层级，0表示不限制
         */
        private int maxDepth = 0;

        /**
         * 最后一层数据是否需要将其子节点合并
         */
        private boolean lastLevelMerge = false;

        /**
         * 原始数据
         */
        private List<T> dataList;

        public GenericsTree<T> build() {
            if (parentFiled == null) {
                throw new RuntimeException("parentFiled 字段必填.");
            }
            if (childFiled == null) {
                throw new RuntimeException("childFiled 字段必填.");
            }
            if (dataList == null) {
                throw new RuntimeException("dataList 字段必填.");
            }
            if (rootList == null || rootLevelCondition == null) {
                throw new RuntimeException("rootList 字段 rootLevelCondition 字段必填一个.");
            }
            GenericsTree<T> genericsTree = new GenericsTree<>();
            genericsTree.builder = this;
            return genericsTree;
        }

        public Builder<T> setParentFiled(Function<T, Object> parentFiled) {
            this.parentFiled = parentFiled;
            return this;
        }

        public Builder<T> setChildFiled(Function<T, Object> childFiled) {
            this.childFiled = childFiled;
            return this;
        }

        public Builder<T> setRootList(List<T> rootList) {
            this.rootList = rootList;
            return this;
        }

        public Builder<T> setRootLevelCondition(Function<T, Boolean> rootLevelCondition) {
            this.rootLevelCondition = rootLevelCondition;
            return this;
        }

        public Builder<T> setSortCondition(Comparator<T> sortCondition) {
            this.sortCondition = sortCondition;
            return this;
        }

        public Builder<T> addSortLevelCondition(Integer depth, Comparator<T> sortCondition) {
            if (depth < 1) {
                throw new RuntimeException("The depth must be greater than 1");
            }
            sortLevelCondition.put(depth, sortCondition);
            return this;
        }

        public Builder<T> setFillChildNodeMethod(BiConsumer<T, List<T>> fillChildNodeMethod) {
            this.fillChildNodeMethod = fillChildNodeMethod;
            return this;
        }

        public Builder<T> setMaxDepth(int maxDepth) {
            this.maxDepth = maxDepth;
            return this;
        }

        public Builder<T> setLastLevelMerge(boolean lastLevelMerge) {
            this.lastLevelMerge = lastLevelMerge;
            return this;
        }

        public Builder<T> setDataList(List<T> dataList) {
            this.dataList = dataList;
            return this;
        }
    }
}