package com.njzscloud.common.core.tree;


import cn.hutool.core.collection.CollUtil;
import com.njzscloud.common.core.utils.GroupUtil;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树型结构工具
 */
public class Tree {

    /**
     * 创建树形数据(倒序)
     *
     * @param src    源集合
     * @param rootId 根节点 ID
     * @param <T>    节点 ID 类型
     * @return List&lt;? extends TreeNode&lt;T&gt;&gt;
     */
    public static <T> List<? extends TreeNode<T>> listToTreeDesc(List<? extends TreeNode<T>> src, T rootId) {
        return listToTree(src.stream().sorted(Comparator.reverseOrder()).collect(Collectors.toList()), TreeNode::getId, TreeNode::getPid, TreeNode::setChildren, rootId, true);
    }

    /**
     * 创建树形数据(正序)
     *
     * @param src    源集合
     * @param rootId 根节点 ID
     * @param <T>    节点 ID 类型
     * @return List&lt;? extends TreeNode&lt;T&gt;&gt;
     */
    public static <T> List<? extends TreeNode<T>> listToTreeAsc(List<? extends TreeNode<T>> src, T rootId) {
        return listToTree(src.stream().sorted().collect(Collectors.toList()), TreeNode::getId, TreeNode::getPid, TreeNode::setChildren, rootId, false);
    }

    /**
     * 创建树形数据(排序)
     *
     * @param src           源集合
     * @param idFn          id 提供函数
     * @param pidFn         pid 提供函数
     * @param setChildrenFn 子节点设置函数
     * @param rootId        根节点 ID
     * @param reverse       是否倒序
     * @param <M>           源集合元素类型
     * @param <T>           节点 ID 类型
     * @return List&lt;M&gt;
     */
    public static <M extends Comparable<? super M>, T> List<M> listToTree(List<M> src, Function<M, T> idFn, Function<M, T> pidFn, BiConsumer<M, List<M>> setChildrenFn, T rootId, boolean reverse) {
        if (CollUtil.isEmpty(src)) return Collections.emptyList();
        if (reverse) {
            src = src
                    .stream()
                    .sorted(Comparator.reverseOrder())
                    .collect(Collectors.toList());
        }

        Map<T, List<M>> pid_treeNode_map = GroupUtil.k_a(src, pidFn);

        for (M m : src) {
            setChildrenFn.accept(m, pid_treeNode_map.get(idFn.apply(m)));
        }
        return CollUtil.emptyIfNull(pid_treeNode_map.get(rootId));
    }

    /**
     * 创建树形数据(不排序)
     *
     * @param src           源集合
     * @param idFn          id 提供函数
     * @param pidFn         pid 提供函数
     * @param setChildrenFn 子节点设置函数
     * @param rootId        根节点 ID
     * @param <M>           源集合元素类型
     * @param <T>           节点 ID 类型
     * @return List&lt;M&gt;
     */
    public static <M, T> List<M> listToTree(List<M> src, Function<M, T> idFn, Function<M, T> pidFn, BiConsumer<M, List<M>> setChildrenFn, T rootId) {
        if (CollUtil.isEmpty(src)) return Collections.emptyList();

        Map<T, List<M>> pid_treeNode_map = GroupUtil.k_a(src, pidFn);

        for (M m : src) {
            setChildrenFn.accept(m, pid_treeNode_map.get(idFn.apply(m)));
        }
        return CollUtil.emptyIfNull(pid_treeNode_map.get(rootId));
    }

    /**
     * 扁平化树
     *
     * @param src           源集合 树形结构
     * @param getChildrenFn 子节点获取函数
     * @param convert       转换函数, 源集合 元素 -> 结果集合 元素
     * @param <S>           源集合 元素类型
     * @param <D>           结果集合 元素类型
     * @return List&lt;D&gt;
     */
    public static <S, D> List<D> treeToList(List<S> src, Function<S, List<S>> getChildrenFn, Function<S, D> convert) {
        List<D> dest;
        if (CollUtil.isNotEmpty(src)) {
            dest = new ArrayList<>();
            treeToList(src, dest, getChildrenFn, convert);
        } else {
            dest = Collections.emptyList();
        }
        return dest;
    }

    /**
     * 扁平化树
     *
     * @param src           源集合 树形结构
     * @param dest          结果集合
     * @param getChildrenFn 子节点获取函数
     * @param convert       转换函数, 源集合 元素 -> 结果集合 元素
     * @param <S>           源集合 元素类型
     * @param <D>           结果集合 元素类型
     */
    public static <S, D> void treeToList(List<S> src, List<D> dest, Function<S, List<S>> getChildrenFn, Function<S, D> convert) {
        if (CollUtil.isNotEmpty(src)) {
            for (S s : src) {
                D d = convert.apply(s);
                dest.add(d);
                List<S> children = getChildrenFn.apply(s);
                if (CollUtil.isNotEmpty(children)) {
                    Tree.treeToList(children, dest, getChildrenFn, convert);
                }
            }
        }
    }
}
