package org.matrix.util;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.ToString;
import org.testng.Assert;

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

/**
 * TreeUtils.
 *
 * @author Matrix <xhyrzldf@gmail.com>
 * @since 2022/3/16 at 8:24 PM
 * Suffering is the most powerful teacher of life.
 */
public class TreeUtils {
    Map<Long, OrgNode> orgNodeMap = new HashMap<>();

    {
        //初始化测试用的数据
        orgNodeMap.put(1L, new OrgNode(1, "Jack", 4));
        orgNodeMap.put(2L, new OrgNode(2, "Rose", 4));
        orgNodeMap.put(3L, new OrgNode(3, "Lucy", 0));
        orgNodeMap.put(4L, new OrgNode(4, "Alice", 3));
        orgNodeMap.put(5L, new OrgNode(5, "Tom", 3));
        orgNodeMap.put(6L, new OrgNode(6, "Jason", 8));
        orgNodeMap.put(7L, new OrgNode(7, "Emma", 5));
        orgNodeMap.put(8L, new OrgNode(8, "Brown", 5));
    }

    public Optional<OrgNode> getParentOrgNode(long parentId) {
        return Optional.ofNullable(orgNodeMap.get(parentId));
    }

    public void parseTree() {
        //1.从DB中查询出List<OrgNode>
        List<OrgNode> orgNodeList = new ArrayList<>(orgNodeMap.values());
        //2.将List<OrgNode>解析并转化为Node<OrgNodeVo>对象
        List<OrgNodeVo> rootNodeList = parseTreeFromDown(orgNodeList,
                OrgNode::getId,
                orgNode -> getParentOrgNode(orgNode.getParentId()),
                OrgNode::toVo,
                OrgNodeVo::addChildNode);
        Assert.assertTrue(rootNodeList.size() >= 1);
        //打印出结果
        for (OrgNodeVo root : rootNodeList) {
            printTreeNode(root, OrgNodeVo::getName, OrgNodeVo::getChildNodes, 0);
        }
    }

    

    /**
     * 自底向上整理出树结构，并存放在Node对象中
     * 利用Map将原始List的对象的位置重新整理，并将所有的指针都设置好
     *
     * @param originList      原始待整理的列表
     * @param getId           列表里每个对象获取自身id的方法，用于避免重复计算
     * @param getParent       列表里每个对象获取父对象的方法
     * @param mapToTargetNode 列表里的原始对象转化为最终对象的方法，默认为t->t，即不变(传入null即可)
     * @param targetSetChild  用于最终对象设置子节点的方法，形如R.setChild(R),或者R.addChildList(R)
     * @param <R>             最终对象
     * @param <V>             原始对象
     * @param <T>             主键类型
     * @return 根节点对象集合
     */
    public static  <R, V, T extends Number> List<R> parseTreeFromDown(List<V> originList,
                                                              Function<V, T> getId,
                                                              Function<V, Optional<V>> getParent,
                                                              Function<V, R> mapToTargetNode,
                                                              BiConsumer<R, R> targetSetChild) {
        //K为主键id , Value 为最终对象
        Map<T, R> map = new HashMap<>(32);
        List<T> rootIds = new ArrayList<>();
        for (V originNode : originList) {
            //对于所有节点，如果已经遍历过了则直接取已经设置过子节点的引用
            R targetNode = map.getOrDefault(getId.apply(originNode), mapToTargetNode.apply(originNode));
            Optional<V> parentNode = getParent.apply(originNode);
            //查询父节点，如果不存在父节点则证明该节点为根节点，直接放入map中
            if (parentNode.isPresent()) {
                //否则查询该父节点是否已经已经被连接过指针，如果连接过则取出连接过的继续连接，否则进行第一次连接并存入map
                V parent = parentNode.get();
                T parentId = getId.apply(parent);
                R parentInMap = map.get(parentId);
                if (parentInMap == null) {
                    R newNode = mapToTargetNode.apply(parent);
                    targetSetChild.accept(newNode, targetNode);
                    map.put(parentId, newNode);
                } else {
                    targetSetChild.accept(parentInMap, targetNode);
                }
                //连接完处理之后还需要将自身这个节点存入map
                map.put(getId.apply(originNode), targetNode);
            } else {
                //root node
                map.put(getId.apply(originNode), targetNode);
                rootIds.add(getId.apply(originNode));
            }
        }
        //根据rootIds返回所有的顶层节点
        return rootIds.stream().map(map::get).collect(Collectors.toList());
    }

    /**
     * 打印出树状对象结构
     *
     * @param rootNode      树的根节点
     * @param printProperty 想打印的属性
     * @param getChild      子节点
     * @param stackDepth    栈深度
     * @param <T>           节点对象类型
     * @param <R>           要打印的属性类型
     */
    public <T, R> void printTreeNode(T rootNode,
                                     Function<T, R> printProperty,
                                     Function<T, List<T>> getChild,
                                     int stackDepth) {
        StringJoiner joiner = new StringJoiner("\t");
        IntStream.range(0, stackDepth).mapToObj(i -> "").forEach(joiner::add);
        joiner.add(printProperty.apply(rootNode).toString());
        for (T childNode : getChild.apply(rootNode)) {
            int currentLevel = stackDepth + 1;
            printTreeNode(childNode, printProperty, getChild, currentLevel);
        }
    }

    /**
     * data struct
     */
    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    class OrgNode {
        public long id;
        public String name;
        public long parentId;

        public OrgNodeVo toVo() {
            return new OrgNodeVo(id, name, new ArrayList<>(4));
        }
    }

    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    @ToString(exclude = "childNodes")
    class OrgNodeVo {
        public long id;
        public String name;
        public List<OrgNodeVo> childNodes = new ArrayList<>(4);

        public void addChildNode(OrgNodeVo orgNodeVo) {
            childNodes.add(orgNodeVo);
        }
    }
}
