/*
 * Copyright 2015, 2024 StreamEx contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package one.util.streamex;

import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.stream.StreamSupport;

/**
 * @author Tagir Valeev
 */
/* package */abstract class CrossSpliterator<T, A> implements Spliterator<A> {
    long est;
    int splitPos;
    final Spliterator<T>[] spliterators;
    final Collection<T>[] collections;

    @SuppressWarnings("unchecked")
    CrossSpliterator(Collection<? extends Collection<T>> source) {
        this.splitPos = 0;
        long est = 1;
        try {
            for (Collection<T> c : source) {
                long size = c.size();
                est = StrictMath.multiplyExact(est, size);
            }
        } catch (ArithmeticException e) {
            est = Long.MAX_VALUE;
        }
        this.est = est;
        this.collections = source.toArray(new Collection[0]);
        this.spliterators = new Spliterator[collections.length];
    }

    abstract Spliterator<A> doSplit(long prefixEst, Spliterator<T>[] prefixSpliterators,
            Collection<T>[] prefixCollections);

    abstract void accumulate(int pos, T t);

    CrossSpliterator(long est, int splitPos, Spliterator<T>[] spliterators, Collection<T>[] collections) {
        this.est = est;
        this.splitPos = splitPos;
        this.spliterators = spliterators;
        this.collections = collections;
    }

    static final class Reducing<T, A> extends CrossSpliterator<T, A> {
        private A[] elements;
        private final BiFunction<A, ? super T, A> accumulator;

        @SuppressWarnings("unchecked")
        Reducing(Collection<? extends Collection<T>> source, A identity, BiFunction<A, ? super T, A> accumulator) {
            super(source);
            this.accumulator = accumulator;
            this.elements = (A[]) new Object[collections.length + 1];
            this.elements[0] = identity;
        }

        private Reducing(long est, int splitPos, BiFunction<A, ? super T, A> accumulator,
                Spliterator<T>[] spliterators, Collection<T>[] collections, A[] elements) {
            super(est, splitPos, spliterators, collections);
            this.accumulator = accumulator;
            this.elements = elements;
        }

        @Override
        public boolean tryAdvance(Consumer<? super A> action) {
            if (elements == null)
                return false;
            if (est < Long.MAX_VALUE && est > 0)
                est--;
            int l = collections.length;
            if (advance(l - 1)) {
                action.accept(elements[l]);
                return true;
            }
            elements = null;
            est = 0;
            return false;
        }

        @Override
        public void forEachRemaining(Consumer<? super A> action) {
            if (elements == null)
                return;
            int l = collections.length;
            A[] e = elements;
            while (advance(l - 1)) {
                action.accept(e[l]);
            }
            elements = null;
            est = 0;
        }

        @Override
        Spliterator<A> doSplit(long prefixEst, Spliterator<T>[] prefixSpliterators, Collection<T>[] prefixCollections) {
            return new Reducing<>(prefixEst, splitPos, accumulator, prefixSpliterators, prefixCollections, elements
                    .clone());
        }

        @Override
        void accumulate(int pos, T t) {
            elements[pos + 1] = accumulator.apply(elements[pos], t);
        }
    }

    static final class ToList<T> extends CrossSpliterator<T, List<T>> {
        private List<T> elements;

        @SuppressWarnings("unchecked")
        ToList(Collection<? extends Collection<T>> source) {
            super(source);
            this.elements = (List<T>) Arrays.asList(new Object[collections.length]);
        }

        private ToList(long est, int splitPos, Spliterator<T>[] spliterators, Collection<T>[] collections,
                List<T> elements) {
            super(est, splitPos, spliterators, collections);
            this.elements = elements;
        }

        @Override
        public boolean tryAdvance(Consumer<? super List<T>> action) {
            if (elements == null)
                return false;
            if (est < Long.MAX_VALUE && est > 0)
                est--;
            if (advance(collections.length - 1)) {
                action.accept(new ArrayList<>(elements));
                return true;
            }
            elements = null;
            est = 0;
            return false;
        }

        @Override
        public void forEachRemaining(Consumer<? super List<T>> action) {
            if (elements == null)
                return;
            List<T> e = elements;
            int l = collections.length - 1;
            while (advance(l)) {
                action.accept(new ArrayList<>(e));
            }
            elements = null;
            est = 0;
        }

        @Override
        Spliterator<List<T>> doSplit(long prefixEst, Spliterator<T>[] prefixSpliterators,
                Collection<T>[] prefixCollections) {
            @SuppressWarnings("unchecked")
            List<T> prefixElements = (List<T>) Arrays.asList(elements.toArray());
            return new ToList<>(prefixEst, splitPos, prefixSpliterators, prefixCollections, prefixElements);
        }

        @Override
        void accumulate(int pos, T t) {
            elements.set(pos, t);
        }
    }

    boolean advance(int i) {
        if (spliterators[i] == null) {
            if (i > 0 && collections[i - 1] != null && !advance(i - 1))
                return false;
            spliterators[i] = collections[i].spliterator();
        }
        Consumer<? super T> action = t -> accumulate(i, t);
        if (!spliterators[i].tryAdvance(action)) {
            if (i == 0 || collections[i - 1] == null || !advance(i - 1))
                return false;
            spliterators[i] = collections[i].spliterator();
            return spliterators[i].tryAdvance(action);
        }
        return true;
    }

    @Override
    public Spliterator<A> trySplit() {
        if (spliterators[splitPos] == null)
            spliterators[splitPos] = collections[splitPos].spliterator();
        Spliterator<T> res = spliterators[splitPos].trySplit();
        if (res == null) {
            if (splitPos == spliterators.length - 1)
                return null;
            @SuppressWarnings("unchecked")
            T[] arr = (T[]) StreamSupport.stream(spliterators[splitPos], false).toArray();
            if (arr.length == 0)
                return null;
            if (arr.length == 1) {
                accumulate(splitPos, arr[0]);
                splitPos++;
                return trySplit();
            }
            spliterators[splitPos] = Spliterators.spliterator(arr, Spliterator.ORDERED);
            return trySplit();
        }
        long prefixEst = Long.MAX_VALUE;
        long newEst = spliterators[splitPos].getExactSizeIfKnown();
        if (newEst == -1) {
            newEst = Long.MAX_VALUE;
        } else {
            try {
                for (int i = splitPos + 1; i < collections.length; i++) {
                    long size = collections[i].size();
                    newEst = StrictMath.multiplyExact(newEst, size);
                }
                if (est != Long.MAX_VALUE)
                    prefixEst = est - newEst;
            } catch (ArithmeticException e) {
                newEst = Long.MAX_VALUE;
            }
        }
        Spliterator<T>[] prefixSpliterators = spliterators.clone();
        Collection<T>[] prefixCollections = collections.clone();
        prefixSpliterators[splitPos] = res;
        this.est = newEst;
        Arrays.fill(spliterators, splitPos + 1, spliterators.length, null);
        return doSplit(prefixEst, prefixSpliterators, prefixCollections);
    }

    @Override
    public long estimateSize() {
        return est;
    }

    @Override
    public int characteristics() {
        int sized = est < Long.MAX_VALUE ? SIZED : 0;
        return ORDERED | sized;
    }
}
