/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */

package org.elasticsearch.common.collect;

import org.elasticsearch.common.Randomness;
import org.elasticsearch.test.ESTestCase;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

public class IteratorsTests extends ESTestCase {
    public void testConcatentation() {
        List<Integer> threeTwoOne = Arrays.asList(3, 2, 1);
        List<Integer> fourFiveSix = Arrays.asList(4, 5, 6);
        Iterator<Integer> concat = Iterators.concat(threeTwoOne.iterator(), fourFiveSix.iterator());
        assertContainsInOrder(concat, 3, 2, 1, 4, 5, 6);
    }

    public void testNoConcatenation() {
        Iterator<Integer> iterator = Iterators.<Integer>concat();
        assertEmptyIterator(iterator);
    }

    public void testEmptyConcatenation() {
        Iterator<Integer> iterator = Iterators.<Integer>concat(empty());
        assertEmptyIterator(iterator);
    }

    public void testMultipleEmptyConcatenation() {
        Iterator<Integer> iterator = Iterators.concat(empty(), empty());
        assertEmptyIterator(iterator);
    }

    public void testSingleton() {
        int value = randomInt();
        assertSingleton(value, singletonIterator(value));
    }

    public void testEmptyBeforeSingleton() {
        int value = randomInt();
        assertSingleton(value, empty(), singletonIterator(value));
    }

    public void testEmptyAfterSingleton() {
        int value = randomInt();
        assertSingleton(value, singletonIterator(value), empty());
    }

    public void testRandomSingleton() {
        int numberOfIterators = randomIntBetween(1, 1000);
        int singletonIndex = randomIntBetween(0, numberOfIterators - 1);
        int value = randomInt();
        @SuppressWarnings({ "rawtypes", "unchecked" })
        Iterator<Integer>[] iterators = new Iterator[numberOfIterators];
        for (int i = 0; i < numberOfIterators; i++) {
            iterators[i] = i != singletonIndex ? empty() : singletonIterator(value);
        }
        assertSingleton(value, iterators);
    }

    public void testRandomIterators() {
        int numberOfIterators = randomIntBetween(1, 1000);
        @SuppressWarnings({ "rawtypes", "unchecked" })
        Iterator<Integer>[] iterators = new Iterator[numberOfIterators];
        List<Integer> values = new ArrayList<>();
        for (int i = 0; i < numberOfIterators; i++) {
            int numberOfValues = randomIntBetween(0, 256);
            List<Integer> theseValues = new ArrayList<>();
            for (int j = 0; j < numberOfValues; j++) {
                int value = randomInt();
                values.add(value);
                theseValues.add(value);
            }
            iterators[i] = theseValues.iterator();
        }
        assertContainsInOrder(Iterators.concat(iterators), values.toArray(new Integer[values.size()]));
    }

    public void testTwoEntries() {
        int first = randomInt();
        int second = randomInt();
        Iterator<Integer> concat = Iterators.concat(singletonIterator(first), empty(), empty(), singletonIterator(second));
        assertContainsInOrder(concat, first, second);
    }

    public void testNull() {
        try {
            Iterators.concat((Iterator<?>) null);
            fail("expected " + NullPointerException.class.getSimpleName());
        } catch (NullPointerException e) {

        }
    }

    public void testNullIterator() {
        try {
            Iterators.concat(singletonIterator(1), empty(), null, empty(), singletonIterator(2));
            fail("expected " + NullPointerException.class.getSimpleName());
        } catch (NullPointerException e) {

        }
    }

    public void testArrayIterator() {
        Integer[] array = randomIntegerArray();
        Iterator<Integer> iterator = Iterators.forArray(array);

        int i = 0;
        while (iterator.hasNext()) {
            assertEquals(array[i++], iterator.next());
        }
        assertEquals(array.length, i);
    }

    public void testArrayIteratorForEachRemaining() {
        Integer[] array = randomIntegerArray();
        Iterator<Integer> iterator = Iterators.forArray(array);

        AtomicInteger index = new AtomicInteger();
        iterator.forEachRemaining(i -> assertEquals(array[index.getAndIncrement()], i));
        assertEquals(array.length, index.get());
    }

    public void testArrayIteratorIsUnmodifiable() {
        Integer[] array = randomIntegerArray();
        Iterator<Integer> iterator = Iterators.forArray(array);

        expectThrows(UnsupportedOperationException.class, iterator::remove);
    }

    public void testArrayIteratorThrowsNoSuchElementExceptionWhenDepleted() {
        Integer[] array = randomIntegerArray();
        Iterator<Integer> iterator = Iterators.forArray(array);
        for (int i = 0; i < array.length; i++) {
            iterator.next();
        }

        expectThrows(NoSuchElementException.class, iterator::next);
    }

    public void testArrayIteratorOnNull() {
        expectThrows(NullPointerException.class, "Unable to iterate over a null array", () -> Iterators.forArray(null));
    }

    public void testFlatMap() {
        final var array = randomIntegerArray();
        assertEmptyIterator(Iterators.flatMap(Iterators.forArray(array), i -> Iterators.concat()));
        assertEmptyIterator(Iterators.flatMap(Iterators.concat(), i -> Iterators.single("foo")));

        final var index = new AtomicInteger();
        Iterators.flatMap(Iterators.forArray(array), Iterators::single)
            .forEachRemaining(i -> assertEquals(array[index.getAndIncrement()], i));
        assertEquals(array.length, index.get());

        index.set(0);
        Iterators.flatMap(Iterators.forArray(array), i -> Iterators.forArray(new Integer[] { i, i, i }))
            .forEachRemaining(i -> assertEquals(array[(index.getAndIncrement() / 3)], i));
        assertEquals(array.length * 3, index.get());

        final var input = new ArrayList<>();
        for (int i = 1; i <= 4; i++) {
            IntStream.range(0, between(0, 2)).forEachOrdered(ignored -> input.add(0));
            input.add(i);
            IntStream.range(0, between(0, 2)).forEachOrdered(ignored -> input.add(0));
        }

        index.set(0);
        final var expectedArray = new Integer[] { 0, 0, 1, 0, 1, 2, 0, 1, 2, 3 };
        Iterators.flatMap(input.listIterator(), i -> IntStream.range(0, (int) i).iterator())
            .forEachRemaining(i -> assertEquals(expectedArray[(index.getAndIncrement())], i));
        assertEquals(expectedArray.length, index.get());
    }

    private static Integer[] randomIntegerArray() {
        return Randomness.get().ints(randomIntBetween(0, 1000)).boxed().toArray(Integer[]::new);
    }

    private <T> Iterator<T> singletonIterator(T value) {
        return Collections.singleton(value).iterator();
    }

    @SafeVarargs
    @SuppressWarnings({ "unchecked", "varargs" })
    private <T> void assertSingleton(T value, Iterator<T>... iterators) {
        Iterator<T> concat = Iterators.concat(iterators);
        assertContainsInOrder(concat, value);
    }

    private <T> Iterator<T> empty() {
        return new Iterator<T>() {
            @Override
            public boolean hasNext() {
                return false;
            }

            @Override
            public T next() {
                throw new NoSuchElementException();
            }
        };
    }

    @SafeVarargs
    @SuppressWarnings({ "varargs" })
    private <T> void assertContainsInOrder(Iterator<T> iterator, T... values) {
        for (T value : values) {
            assertTrue(iterator.hasNext());
            assertEquals(value, iterator.next());
        }
        assertNoSuchElementException(iterator);
    }

    private <T> void assertEmptyIterator(Iterator<T> iterator) {
        assertFalse(iterator.hasNext());
        assertNoSuchElementException(iterator);
    }

    private <T> void assertNoSuchElementException(Iterator<T> iterator) {
        try {
            iterator.next();
            fail("expected " + NoSuchElementException.class.getSimpleName());
        } catch (NoSuchElementException e) {

        }
    }
}
