/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.jena.atlas.data;

import java.io.*;
import java.util.*;

import org.apache.jena.atlas.AtlasException;
import org.apache.jena.atlas.data.AbortableComparator.Finish;
import org.apache.jena.atlas.iterator.Iter;
import org.apache.jena.atlas.iterator.IteratorCloseable;
import org.apache.jena.atlas.lib.Sink;

/**
 * <p>
 * This data bag will gather items in memory until a size threshold is passed,
 * at which point it will write out all of the items to disk using the supplied
 * serializer.
 * </p>
 * <p>
 * After adding is finished, call {@link #iterator()} to set up the data bag for
 * reading back items and iterating over them. The iterator will retrieve the
 * items in sorted order using the supplied comparator.
 * </p>
 * <p>
 * IMPORTANT: You may not add any more items after this call. You may
 * subsequently call {@link #iterator()} multiple times which will give you a
 * new iterator for each invocation. If you do not consume the entire iterator,
 * you should call {@link Iter#close(Iterator)} to close any FileInputStreams
 * associated with the iterator.
 * </p>
 * <p>
 * Additionally, make sure to call {@link #close()} when you are finished to
 * free any system resources (preferably in a finally block).
 * </p>
 * <p>
 * Implementation Notes: Data is stored in an ArrayList as it comes in. When it
 * is time to spill, that data is sorted and written to disk. An iterator will
 * read in each file and perform a merge-sort as the results are returned.
 * </p>
 */
public class SortedDataBag<E> extends AbstractDataBag<E> {
    /**
     * The the maximum number of files to merge at the same time. Without this, you
     * can run out of file handles and other bad things.
     */
    protected static int MAX_SPILL_FILES = 100;

    protected final ThresholdPolicy<E> policy;
    protected final SerializationFactory<E> serializationFactory;
    protected final AbortableComparator<E> comparator;

    protected boolean finishedAdding = false;
    protected boolean spilled = false;
    protected boolean closed = false;

    public SortedDataBag(ThresholdPolicy<E> policy, SerializationFactory<E> serializerFactory, Comparator<? super E> comparator) {
        this.policy = policy;
        this.serializationFactory = serializerFactory;
        this.comparator = new AbortableComparator<>(comparator);
    }

    /**
     * cancel arranges that further comparisons using the supplied comparator will
     * abandon the sort in progress.
     */
    public void cancel() {
        comparator.cancel();
        close();
    }

    /**
     * isCancelled is true iff cancel has been called on this bags comparator. (Used
     * in testing.)
     */
    public boolean isCancelled() {
        return comparator.cancelled;
    }

    /**
     * isClosed returns true iff this bag has been closed. (Used in testing.)
     */
    public boolean isClosed() {
        return closed;
    }

    protected void checkClosed() {
        if ( closed )
            throw new AtlasException("SortedDataBag is closed, no operations can be performed on it.");
    }

    @Override
    public boolean isSorted() {
        return true;
    }

    @Override
    public boolean isDistinct() {
        return false;
    }

    @Override
    public void add(E item) {
        checkClosed();
        if ( finishedAdding )
            throw new AtlasException("SortedDataBag: Cannot add any more items after the writing phase is complete.");

        if ( policy.isThresholdExceeded() ) {
            spill();
        }

        if ( memory.add(item) ) {
            policy.increment(item);
            size++;
        }
    }

    @SuppressWarnings({"unchecked"})
    protected void spill() {
        // Make sure we have something to spill.
        if ( memory.size() > 0 ) {
            OutputStream out;
            try {
                out = getSpillStream();
            } catch (IOException e) {
                throw new AtlasException(e);
            }

            // Sort the tuples as an array. The CanAbortComparator will sort
            // the array using Arrays.sort. The cast to E[] is safe. If the sort is
            // aborted, don't bother messing around with the serialisation.
            // We'll never get around to using it anyway.

            E[] array = (E[])memory.toArray();
            if ( comparator.abortableSort(array) == Finish.COMPLETED ) {
                Sink<E> serializer = serializationFactory.createSerializer(out);
                try {
                    for ( Object tuple : array ) {
                        serializer.send((E)tuple);
                    }
                }
                finally {
                    serializer.close();
                }
            }

            spilled = true;
            policy.reset();
            memory.clear();
        }
    }

    @Override
    public void flush() {
        spill();
    }

    protected Iterator<E> getInputIterator(File spillFile) throws FileNotFoundException {
        InputStream in = getInputStream(spillFile);
        Iterator<E> deserializer = serializationFactory.createDeserializer(in);
        return Iter.onCloseIO(deserializer, in);
    }

    /**
     * Returns an iterator over a set of elements of type E. If you do not exhaust
     * the iterator, you should call
     * {@link org.apache.jena.atlas.iterator.Iter#close(Iterator)} to be sure any
     * open file handles are closed.
     *
     * @return an Iterator
     */
    @Override
    public Iterator<E> iterator() {
        preMerge();

        return iterator(getSpillFiles().size());
    }

    @SuppressWarnings({"unchecked"})
    private Iterator<E> iterator(int size) {
        checkClosed();

        int memSize = memory.size();

        // Constructing an iterator from this class is not thread-safe (just
        // like all the the other methods)
        if ( !finishedAdding && memSize > 1 ) {
            E[] array = (E[])memory.toArray();
            comparator.abortableSort(array); // don't care if we aborted or not
            memory = Arrays.asList(array);
        }

        finishedAdding = true;

        if ( spilled ) {
            List<Iterator<E>> inputs = new ArrayList<>(size + (memSize > 0 ? 1 : 0));

            if ( memSize > 0 ) {
                inputs.add(memory.iterator());
            }

            for ( int i = 0 ; i < size ; i++ ) {
                File spillFile = getSpillFiles().get(i);
                try {
                    Iterator<E> irc = getInputIterator(spillFile);
                    inputs.add(irc);
                } catch (FileNotFoundException e) {
                    // Close any open streams before we throw an exception
                    for ( Iterator<E> it : inputs ) {
                        Iter.close(it);
                    }

                    throw new AtlasException("Cannot find one of the spill files", e);
                }
            }

            SpillSortIterator<E> ssi = new SpillSortIterator<>(inputs, comparator);
            registerCloseableIterator(ssi);

            return ssi;
        } else {
            if ( memSize > 0 ) {
                return memory.iterator();
            } else {
                return Iter.nullIterator();
            }
        }
    }

    private void preMerge() {
        if ( getSpillFiles() == null || getSpillFiles().size() <= MAX_SPILL_FILES ) {
            return;
        }

        try {
            while (getSpillFiles().size() > MAX_SPILL_FILES) {
                Sink<E> sink = serializationFactory.createSerializer(getSpillStream());
                Iterator<E> ssi = iterator(MAX_SPILL_FILES);
                try {
                    while (ssi.hasNext()) {
                        sink.send(ssi.next());
                    }
                }
                finally {
                    Iter.close(ssi);
                    sink.close();
                }

                List<File> toRemove = new ArrayList<>(MAX_SPILL_FILES);
                for ( int i = 0 ; i < MAX_SPILL_FILES ; i++ ) {
                    File file = getSpillFiles().get(i);
                    file.delete();
                    toRemove.add(file);
                }

                getSpillFiles().removeAll(toRemove);

                memory = new ArrayList<>();
            }
        } catch (IOException e) {
            throw new AtlasException(e);
        }
    }

    @Override
    public void close() {
        if ( !closed ) {
            closeIterators();
            deleteSpillFiles();

            memory = null;
            closed = true;
        }
    }

    /**
     * An iterator that handles getting the next tuple from the bag.
     */
    protected static class SpillSortIterator<T> implements IteratorCloseable<T> {
        private final List<Iterator<T>> inputs;
        private final Comparator<? super T> comp;
        private final PriorityQueue<Item<T>> minHeap;

        public SpillSortIterator(List<Iterator<T>> inputs, Comparator<? super T> comp) {
            this.inputs = inputs;
            this.comp = comp;
            this.minHeap = new PriorityQueue<>(inputs.size());

            // Prime the heap
            for ( int i = 0 ; i < inputs.size() ; i++ ) {
                replaceItem(i);
            }
        }

        private void replaceItem(int index) {
            Iterator<T> it = inputs.get(index);
            if ( it.hasNext() ) {
                T tuple = it.next();
                minHeap.add(new Item<>(index, tuple, comp));
            }
        }

        @Override
        public boolean hasNext() {
            return (minHeap.peek() != null);
        }

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

            Item<T> curr = minHeap.poll();
            // Read replacement item
            replaceItem(curr.getIndex());

            return curr.getTuple();
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("SpillSortIterator.remove");
        }

        @Override
        public void close() {
            for ( Iterator<T> it : inputs ) {
                Iter.close(it);
            }
        }

        private final class Item<U> implements Comparable<Item<U>> {
            private final int index;
            private final U tuple;
            private final Comparator<? super U> c;

            public Item(int index, U tuple, Comparator<? super U> c) {
                this.index = index;
                this.tuple = tuple;
                this.c = c;
            }

            public int getIndex() {
                return index;
            }

            public U getTuple() {
                return tuple;
            }

            @Override
            @SuppressWarnings("unchecked")
            public int compareTo(Item<U> o) {
                return (null != c) ? c.compare(tuple, o.getTuple()) : ((Comparable<U>)tuple).compareTo(o.getTuple());
            }

            @SuppressWarnings("unchecked")
            @Override
            public boolean equals(Object obj) {
                if ( obj instanceof Item ) {
                    return compareTo((Item<U>)obj) == 0;
                }

                return false;
            }

            @Override
            public int hashCode() {
                return tuple.hashCode();
            }
        }

    }
}
