package cn.lihongjie.coal.common;

import cn.lihongjie.coal.exception.BizException;

import io.vavr.Function2;

import lombok.Data;
import lombok.experimental.UtilityClass;

import org.apache.poi.ss.formula.functions.T;

import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;

@UtilityClass
public class TreeUtils {


    public static <T> Iterable<T> buildTreeFromList(Iterable<T> src, Function<T, String> idGetter, Function<T, String> parentIdGetter, Function2<T, T, Void> childrenAppender) {
        Map<String, T> map = new HashMap<>();
        List<T> roots = new ArrayList<>();
        for (T t : src) {
            map.put(idGetter.apply(t), t);
        }
        for (T t : src) {
            String parentId = parentIdGetter.apply(t);
            if (parentId == null || parentId.isEmpty()) {
                roots.add(t);
            } else {
                T parent = map.get(parentId);
                if (parent == null) {
                    throw new BizException("父节点不存在 " + parentId);
                }

                childrenAppender.apply(parent, t);
            }
        }
        return roots;
    }


    public static <T> Iterable<T> allChildren(T src, Function<T, Iterable<T>> childrenGetter) {

        ArrayList<T> ans = new ArrayList<>();

        var seen = new HashSet<>();

        Stack<T> stack = new Stack<>();

        stack.push(src);

        while (!stack.isEmpty()) {
            T pop = stack.pop();
            if (!seen.add(pop)) {
                throw new BizException("循环引用 " + pop.toString());
            }
            if (pop != src) {

                ans.add(pop);
            }
            Iterable<T> children = childrenGetter.apply(pop);
            if (children != null) {

                for (T child : children) {
                    stack.push(child);
                }
            }
        }
        return ans;
    }

    public static <T>void dfsList(Iterable<T> src, Function<T, Iterable<T>> getter, Consumer<DFSCtx<T>> handler,  Function<T, Object> hashFn){

        src.forEach(t -> dfs(t, getter, handler, hashFn));
    }

    public static <T>void dfs(T src, Function<T, Iterable<T>> getter, Consumer<DFSCtx<T>> handler,  Function<T, Object> hashFn){

        var seen = new HashSet<>();

        Stack<DFSCtx<T>> stack = new Stack<>();

        stack.push(new DFSCtx<>(src, null, src));

        while (!stack.isEmpty()) {
            DFSCtx<T> pop = stack.pop();
            if (hashFn!=null && !seen.add(hashFn.apply(pop.current))) {
                throw new BizException("循环引用 " + pop.current.toString());
            }
            handler.accept(pop);

            Iterable<T> children = getter.apply(pop.current);
            if (children != null) {

                for (T child : children) {
                    stack.push(new DFSCtx<>(src, pop.current, child));
                }
            }
        }

    }

    public static <T> void forEachChildren(
            T src,
            Function<T, Iterable<T>> getter,
            Function2<T, Iterable<T>, Void> setter,
            Function<Iterable<T>, Iterable<T>> mapper) {

        var seen = new HashSet<>();

        Stack<T> stack = new Stack<>();

        stack.push(src);

        while (!stack.isEmpty()) {
            T pop = stack.pop();
//            if (!seen.add(pop)) {
//                throw new BizException("循环引用 " + pop.toString());
//            }

            Iterable<T> children = getter.apply(pop);
            if (children != null) {

                for (T child : children) {
                    stack.push(child);
                }
                setter.apply(pop, mapper.apply( children));
            }
        }
    }

    public static <T> List<T> treeToList(Iterable<T> src) {
        return treeToList(src, n -> (Iterable<T>) ReflectUtils.invokeMethod(n, "getChildren"));
    }

    public static <T> List<T> treeToList(Iterable<T> src, Function<T, Iterable<T>> getChildren) {
        return treeToList(src, getChildren, ReflectUtils::getId);
    }

    public static <T> List<T> treeToList(Iterable<T> src, Function<T, Iterable<T>> getChildren, Function<T, Object> hashFn) {
        List<T> ans = new ArrayList<>();
        dfsList(
                src,
                getChildren,
                x -> {
                    ans.add(x.current);
                },
                hashFn);
        return ans;
    }

    @Data
    public static class DFSCtx<T>{

        private T root;
        private T parent;
        private T current;

        public DFSCtx(T root, T parent, T current) {
            this.root = root;
            this.parent = parent;
            this.current = current;
        }
    }
}
