package com.sptg.common.core.util;

import java.util.*;

/**
 * title: TreeUtil
 * package: com.sptg.common.core.util
 * description: 树形工具类
 *
 *                  
 *                 
 */
public class TreeUtil {

    public interface TreeEntity<Entity>{
        /**
         * 获取主键
         * @return Long
         */
        Long getId();

        /**
         * 父级Id
         * @return Long
         */
        Long getPid();

        /**
         * 排序值
         * @return Integer
         */
        Integer getSort();

        /**
         * 设置子节点
         * @param entityList 子节点列表
         */
        void setSubList(List<Entity> entityList);

        /**
         * 获取子节点
         * @return List<Entity>
         */
        List<Entity> getSubList();
    }

    public static <Entity extends TreeEntity<Entity>> List<Entity> tree(List<Entity> entityList){
        if (entityList.size() == 0){
            return entityList;
        }
        entityList.sort(Comparator.comparing(Entity :: getSort));
        Map<Long, List<Entity>> groupMap = new HashMap<>(16);
        entityList.forEach(entity -> {
            List<Entity> list = groupMap.get(entity.getPid());
            if (null == list){
                list = new ArrayList<>();
                list.add(entity);
                groupMap.put(entity.getPid(),list);
            }else {
                list.add(entity);
            }
        });
        //  得到Pid列表
        List<Long> pidList = new ArrayList<>(groupMap.keySet());
        pidList.sort(Comparator.comparing(Long::longValue));
        for (int i = pidList.size() - 1; i > 0; i--){
            //  数值最大的Pid
            long thisKey = pidList.get(i);
            //  寻找父节点
            for (int j = i - 1; j >= 0; j--){
                //  距离最近Pid
                long lastKey = pidList.get(j);
                List<Entity> list = groupMap.get(lastKey);
                //  递归
                if (recursion(list,thisKey,groupMap)){
                    break;
                }
            }
        }
        return groupMap.get(pidList.get(0));
    }

    private static <Entity extends TreeEntity<Entity>> boolean recursion(List<Entity> list,long thisKey,Map<Long, List<Entity>> groupMap){
        for (Entity entity : list){
            if (entity.getId().equals(thisKey)){
                List<Entity> children = groupMap.get(thisKey);
                entity.setSubList(children);
                groupMap.remove(thisKey);
                return true;
            }else if (null != entity.getSubList()){
                return recursion(entity.getSubList(),thisKey,groupMap);
            }
        }
        return false;
    }

}
