package au.id.zancanaro.javacheck;

import java.io.IOException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.IntStream;
import java.util.stream.Stream;

@SuppressWarnings("unused")
public class ShrinkTree<T> {
    private final T value;
    private final Supplier<Stream<ShrinkTree<T>>> children;

    private ShrinkTree(T value, Supplier<Stream<ShrinkTree<T>>> children) {
        this.value = value;
        this.children = children;
    }

    public T getValue() {
        return value;
    }

    public Stream<ShrinkTree<T>> getChildren() {
        return children.get();
    }

    @SuppressWarnings("Convert2MethodRef")
    public static <T> ShrinkTree<T> pure(T value) {
        // Converting the () -> Stream.empty() into Stream::empty actually
        // breaks Java's generic type inference. Who would have thought?
        return new ShrinkTree<>(value, () -> Stream.empty());
    }

    public static <T> ShrinkTree<T> join(ShrinkTree<ShrinkTree<T>> tree) {
        return new ShrinkTree<>(
                tree.getValue().getValue(),
                () -> Stream.concat(
                        tree.getChildren().map(ShrinkTree::join),
                        tree.getValue().getChildren()));
    }

    private static <T> List<T> makeHeadList(ShrinkTree<T>[] trees) {
        List<T> heads = new ArrayList<>(trees.length);
        for (ShrinkTree<T> tree : trees) {
            heads.add(tree.getValue());
        }
        return heads;
    }

    @SuppressWarnings("unchecked")
    public static <T> Stream<ShrinkTree<T>[]> promoteChildren(ShrinkTree<? super T>[] trees) {
        return IntStream.range(0, trees.length)
                .mapToObj(index -> trees[index].getChildren().map(child ->
                        IntStream.range(0, trees.length)
                                .mapToObj(i -> (i == index ? child : trees[i]))
                                .toArray(ShrinkTree[]::new)))
                .flatMap(x -> x)
                .map(x -> (ShrinkTree<T>[]) x);
    }

    public static <T> Stream<ShrinkTree<T>[]> removeChildren(ShrinkTree<? super T>[] trees) {
        return IntStream.range(0, trees.length)
                .mapToObj(index -> IntStream.range(0, trees.length)
                        .filter(i -> i != index)
                        .mapToObj(i -> trees[i])
                        .toArray(ShrinkTree[]::new));
    }

    public static <T> Stream<ShrinkTree<T>[]> removeAndPromoteChildren(ShrinkTree<? super T>[] trees) {
        return Stream.concat(removeChildren(trees), promoteChildren(trees));
    }

    public static <T> ShrinkTree<List<T>> combine(
            ShrinkTree<T>[] trees,
            Function<ShrinkTree<T>[], Stream<ShrinkTree<T>[]>> processChildren) {
        return new ShrinkTree<>(
                makeHeadList(trees),
                () -> processChildren.apply(trees)
                        .map(shrinks -> combine(shrinks, processChildren)));
    }

    public <R> ShrinkTree<R> map(Function<? super T, ? extends R> f) {
        return new ShrinkTree<>(
                f.apply(this.value),
                () -> this.getChildren().map(tree -> tree.map(f)));
    }

    public <R> ShrinkTree<R> flatMap(Function<?super T, ? extends ShrinkTree<R>> f) {
        return ShrinkTree.join(this.map(f));
    }

    public ShrinkTree<T> filter(Predicate<? super T> predicate) {
        if (predicate.test(this.getValue())) {
            return new ShrinkTree<>(
                    this.getValue(),
                    () -> this.getChildren()
                            .filter(tree -> predicate.test(tree.getValue()))
                            .map(tree -> tree.filter(predicate)));
        } else {
            throw new IllegalArgumentException("Current value doesn't match predicate: whoops!");
        }
    }

    public ShrinkTree<T> withShrinkStrategy(ShrinkStrategy<T> strategy) {
        return new ShrinkTree<>(this.getValue(), () -> strategyStream(this.getValue(), strategy));
    }

    private static <T> Stream<ShrinkTree<T>> strategyStream(final T value, final ShrinkStrategy<T> strategy) {
        return strategy.shrink(value)
                .map(v -> new ShrinkTree<T>(v, () -> strategyStream(v, strategy)));
    }

    public void print(Writer output) throws IOException {
        print(output, Object::toString);
    }

    public void print(Writer output, Function<? super T, ? extends String> toString) throws IOException {
        output.write(toString.apply(this.getValue()));
        output.write('[');
        Iterator<ShrinkTree<T>> iterator = children.get().iterator();
        while (iterator.hasNext()) {
            iterator.next().print(output, toString);
            if (iterator.hasNext()) {
                output.write(' ');
            }
        }
        output.write(']');
        output.flush();
    }
}