package com.pashanhoo.common.util;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import org.apache.commons.collections.CollectionUtils;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

public class TreeUtil {

    /**
     * 比较规则
     */
    private static Comparator comparator;

    /**
     * 树节点Class类型
     */
    private static Class<?> treeNodeClass;

    /**
     * 主键的get方法对应的Method对象
     */
    private static Method getKeyMethod;

    /**
     * 获取父节点id的get方法对应的Method对象
     */
    private static Method getParentIdMethod;

    /**
     * 设置子元素对应的方法的Method对象
     */
    private static Method setChildrenMethod;

    /**
     * @param comparator
     * @param treeNodeClass
     * @param key           首字母大写的主键字段名称
     */
    public static void setRule(Comparator comparator, Class<?> treeNodeClass, String key) {
        TreeUtil.comparator = comparator;
        TreeUtil.treeNodeClass = treeNodeClass;
        try {
            TreeUtil.getKeyMethod = treeNodeClass.getMethod("get" + key);
            TreeUtil.getParentIdMethod = treeNodeClass.getMethod("getParentId");
            TreeUtil.setChildrenMethod = treeNodeClass.getMethod("setChildren", List.class);
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
        }

    }

    /**
     * 将符合条件的List转为Tree
     *
     * @param
     * @param rootId
     * @return
     */

    public static <TreeNode> List<TreeNode> listToTree(List<TreeNode> treeNodes, String rootId) {
        if (CollectionUtils.isEmpty(treeNodes)) {
            return Lists.newArrayList();
        }

        //parent -> [data1, data2 ...]
        Multimap<String, TreeNode> parentMap = ArrayListMultimap.create();
        //记录parentMap
        List<TreeNode> rootList = Lists.newArrayList();

        for (TreeNode treeNode : treeNodes) {
            String parentId = null;
            try {
                parentId = (String) getParentIdMethod.invoke(treeNode);
            } catch (IllegalAccessException e) {
                e.printStackTrace();
            } catch (InvocationTargetException e) {
                e.printStackTrace();
            }

            parentMap.put(parentId, treeNode);

            //从根节点的下一级构造树

            //rootId为null，从最顶部构造整棵树
            if (rootId == null) {
                if (parentId == null) {
                    rootList.add(treeNode);
                }
            }
            if (StringUtils.isNotBlank(parentId) && parentId.equals(rootId)) {
                rootList.add(treeNode);
            }
        }

        // 对树节点进行排序
        Collections.sort(rootList, comparator);

        // 递归生成树
        transformTree(rootList, parentMap);
        return rootList;
    }

    private static <TreeNode> void transformTree(List<TreeNode> treeNodeList, Multimap<String, TreeNode> parentMap) {
        for (int i = 0; i < treeNodeList.size(); i++) {
            // 遍历该层的每个元素
            TreeNode treeNode = treeNodeList.get(i);
            //获取当前节点下一层节点的List数据
            List<TreeNode> tempList = null;
            try {
                tempList = (List<TreeNode>) parentMap.get((String) getKeyMethod.invoke(treeNode));
            } catch (IllegalAccessException e) {
                e.printStackTrace();
            } catch (InvocationTargetException e) {
                e.printStackTrace();
            }
            if (CollectionUtils.isNotEmpty(tempList)) {
                //首先排序下一层节点
                Collections.sort(tempList, comparator);
                //将下一层节点数据设置为当前节点的Children数据
                try {
                    setChildrenMethod.invoke(treeNode, tempList);
                } catch (IllegalAccessException e) {
                    e.printStackTrace();
                } catch (InvocationTargetException e) {
                    e.printStackTrace();
                }
                // 进入到下一层处理
                transformTree(tempList, parentMap);
            }
        }

    }
}