package au.id.zancanaro.javacheck;

import au.id.zancanaro.javacheck.object.GeneratorProvider;
import au.id.zancanaro.javacheck.object.ObjectGenerator;

import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

@SuppressWarnings({"unused", "WeakerAccess"})
public final class Generators {
    private Generators() {
    }

    /**
     * Create a generator which explicitly depends on the current "size"
     * parameter.
     *
     * @param makeGenerator A function from size to a generator
     * @param <T>           The static type of the returned generator
     * @return The generator returned by makeGenerator
     */
    public static <T> Generator<T> sized(Function<Integer, Generator<T>> makeGenerator) {
        return (random, size) -> makeGenerator.apply(size).generate(random, size);
    }

    /**
     * Remove a generator's shrink tree.
     *
     * @param gen The generator
     * @param <T> The generator's static type
     * @return A new generator which is the previous generator with shrink tree
     * removed
     */
    public static <T> Generator<T> noShrink(Generator<T> gen) {
        return gen.withShrinkStrategy(value -> Stream.empty());
    }

    @SafeVarargs
    public static <T> Generator<T> oneOf(Generator<? extends T>... gens) {
        return integer(0, gens.length).flatMap(index -> gens[index].map(x -> (T) x));
    }

    @SafeVarargs
    public static <T> Generator<T> elements(T... elements) {
        return elements(Arrays.asList(elements));
    }

    public static <T> Generator<T> elements(List<T> elements) {
        return integer(0, elements.size()).map(elements::get);
    }

    public static Generator<Boolean> bool() {
        return (random, size) ->
                ShrinkTree.pure(random.nextBoolean())
                        .withShrinkStrategy(boolShrinkStrategy());
    }

    private static ShrinkStrategy<Boolean> boolShrinkStrategy() {
        return value -> (value ? Stream.of(false) : Stream.empty());
    }

    public static Generator<Long> longInteger(long lower, long upper) {
        return (random, size) -> {
            long value = random.longs(lower, upper).findFirst().getAsLong();
            long bound = lower > 0 ? lower : (upper < 0 ? upper : 0);
            return ShrinkTree.pure(value)
                    .withShrinkStrategy(longShrinkStrategy(bound));
        };
    }

    public static Generator<Long> longInteger() {
        return sized(size -> longInteger(-size, size));
    }

    private static ShrinkStrategy<Long> longShrinkStrategy(final long bound) {
        return value -> StreamSupport.stream(new Spliterators.AbstractSpliterator<Long>(Long.MAX_VALUE, Spliterator.ORDERED) {
            long curr = value - bound;

            @Override
            public boolean tryAdvance(Consumer<? super Long> action) {
                if (curr == 0) {
                    return false;
                } else {
                    action.accept(value - curr);
                    curr /= 2;
                    return true;
                }
            }
        }, false);
    }

    public static Generator<Integer> integer(int lower, int upper) {
        return (random, size) -> {
            int value = random.ints(lower, upper).findFirst().getAsInt();
            int bound = lower > 0 ? lower : (upper < 0 ? upper : 0);
            return ShrinkTree.pure(value)
                    .withShrinkStrategy(intShrinkStrategy(bound));
        };
    }

    public static Generator<Integer> integer() {
        return sized(size -> integer(-size, size));
    }

    public static Generator<Integer> natural() {
        return sized(size -> integer(0, size));
    }

    private static ShrinkStrategy<Integer> intShrinkStrategy(final int bound) {
        return value -> StreamSupport.stream(new Spliterators.AbstractSpliterator<Integer>(Long.MAX_VALUE, Spliterator.ORDERED) {
            int curr = value - bound;

            @Override
            public boolean tryAdvance(Consumer<? super Integer> action) {
                if (curr == 0) {
                    return false;
                } else {
                    action.accept(value - curr);
                    curr /= 2;
                    return true;
                }
            }
        }, false);
    }

    public static Generator<Double> doublePrecision(double lower, double upper) {
        return (random, size) -> {
            double value = random.doubles(lower, upper).findFirst().getAsDouble();
            double bound = lower > 0 ? lower : (upper < 0 ? upper : 0);
            return ShrinkTree.pure(value)
                    .withShrinkStrategy(doubleShrinkStrategy(bound, Double.MIN_NORMAL /* maybe pick a bigger epsilon? */));
        };
    }

    public static Generator<Double> doublePrecision() {
        return sized(size -> doublePrecision(-size, size));
    }

    private static ShrinkStrategy<Double> doubleShrinkStrategy(final double bound, double epsilon) {
        return value -> StreamSupport.stream(new Spliterators.AbstractSpliterator<Double>(Long.MAX_VALUE, Spliterator.ORDERED) {
            double curr = value - bound;

            @Override
            public boolean tryAdvance(Consumer<? super Double> action) {
                if (Math.abs(curr) < epsilon) {
                    return false;
                } else {
                    action.accept(value - curr);
                    curr /= 2;
                    return true;
                }
            }
        }, false);
    }

    public static <T> Generator<List<T>> listOf(Generator<T> gen, int minElements, int maxElements) {
        return (random, size) -> {
            Generator<Integer> countGen = sized(s -> integer(minElements, maxElements));
            int count = countGen.generate(random, size).getValue();
            return Generator.list(count, gen)
                    .generate(random, size)
                    .filter(list -> minElements <= list.size() && list.size() < maxElements)
                    .map(Collections::unmodifiableList);
        };
    }

    public static <T> Generator<List<T>> listOf(Generator<T> gen) {
        return (random, size) -> {
            Generator<Integer> countGen = sized(s -> integer(0, s));
            int count = countGen.generate(random, size).getValue();
            return Generator.list(count, gen)
                    .generate(random, size)
                    .map(Collections::unmodifiableList);
        };
    }

    @SuppressWarnings("unchecked")
    public static <K, V> Generator<Map<K, V>> mapOf(Generator<K> keyGen, Generator<V> valueGen) {
        return (random, size) -> {
            Generator<Integer> countGen = sized(s -> integer(0, s));
            int count = countGen.generate(random, size).getValue();
            return Generator.list(count, Generator.tuple(keyGen, valueGen))
                    .generate(random, size)
                    .map(pairs -> pairs.stream()
                            .collect(Collectors.toMap(
                                    pair -> (K) pair.get(0),
                                    pair -> (V) pair.get(1),
                                    (first, second) -> second)))
                    .map(Collections::unmodifiableMap);
        };
    }

    public static Generator<Character> character() {
        return integer(0, 256).map(i -> (char) i.intValue());
    }

    public static Generator<Character> asciiCharacter() {
        return integer(32, 127).map(i -> (char) i.intValue());
    }

    public static Generator<Character> alphaNumericCharacter() {
        return oneOf(
                integer(48, 58),
                integer(65, 91),
                integer(97, 123)).map(i -> (char) i.intValue());
    }

    public static Generator<Character> alphaCharacter() {
        return oneOf(
                integer(65, 91),
                integer(97, 123)).map(i -> (char) i.intValue());
    }

    private static String makeString(Character[] arr) {
        StringBuilder builder = new StringBuilder(arr.length);
        for (Character c : arr) {
            builder.append(c);
        }
        return builder.toString();
    }

    public static Generator<String> string() {
        return stringOf(character());
    }

    public static Generator<String> stringOf(Generator<Character> charGen) {
        return listOf(charGen).map(list -> {
            char[] chars = new char[list.size()];
            int i = 0;
            for (Character c : list) {
                chars[i++] = c;
            }
            return String.valueOf(chars);
        });
    }

    public static <T> Generator<T> ofType(Class<T> type) {
        return new ObjectGenerator<>(type);
    }

    public static <T> Generator<T> ofType(Class<T> type, GeneratorProvider provider) {
        return new ObjectGenerator<>(type, provider);
    }
}