content: /*
   * Copyright (C) 2020 ActiveJ LLC.
   *
   * 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 io.activej.datastream.processor.transformer.sort;
  
  import io.activej.async.AsyncAccumulator;
  import io.activej.common.builder.AbstractBuilder;
  import io.activej.datastream.consumer.AbstractStreamConsumer;
  import io.activej.datastream.consumer.StreamConsumer;
  import io.activej.datastream.processor.reducer.StreamReducer;
  import io.activej.datastream.processor.transformer.StreamTransformer;
  import io.activej.datastream.supplier.StreamDataAcceptor;
  import io.activej.datastream.supplier.StreamSupplier;
  import io.activej.datastream.supplier.StreamSuppliers;
  import io.activej.promise.Promise;
  import io.activej.reactor.ImplicitlyReactive;
  import org.slf4j.Logger;
  
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.Iterator;
  import java.util.List;
  import java.util.concurrent.Executor;
  import java.util.function.Function;
  
  import static io.activej.datastream.processor.reducer.Reducers.deduplicateReducer;
  import static io.activej.datastream.processor.reducer.Reducers.mergeReducer;
  import static org.slf4j.LoggerFactory.getLogger;
  
  /**
   * Represent {@link StreamTransformer} which receives data and saves it in
   * collection, when it receives end of stream it sorts it and streams to destination.
   *
   * @param <K> type of keys
   * @param <T> type of objects
   */
  public final class StreamSorter<K, T> extends ImplicitlyReactive implements StreamTransformer<T, T> {
  	private static final Logger logger = getLogger(StreamSorter.class);
  	private final AsyncAccumulator<? extends List<Integer>> temporaryStreamsAccumulator;
  	private final IStreamSorterStorage<T> storage;
  	private final Function<T, K> keyFunction;
  	private final Comparator<K> keyComparator;
  	private final Comparator<T> itemComparator;
  	private final boolean distinct;
  	private final int itemsInMemory;
  
  	private final Input input;
  	private final StreamSupplier<T> output;
  
  	private Executor sortingExecutor = Runnable::run;
  
  	private StreamSorter(
  		IStreamSorterStorage<T> storage, Function<T, K> keyFunction, Comparator<K> keyComparator, boolean deduplicate,
  		int itemsInMemory
  	) {
  		this.storage = storage;
  		this.keyFunction = keyFunction;
  		this.keyComparator = keyComparator;
  		this.itemComparator = (item1, item2) -> {
  			K key1 = keyFunction.apply(item1);
  			K key2 = keyFunction.apply(item2);
  			return keyComparator.compare(key1, key2);
  		};
  		this.distinct = deduplicate;
  		this.itemsInMemory = itemsInMemory;
  
  		List<Integer> partitionIds = new ArrayList<>();
  		this.input = new Input(partitionIds);
  		this.output = StreamSuppliers.ofPromise(
  			(this.temporaryStreamsAccumulator = AsyncAccumulator.create(partitionIds))
  				.get()
  				.then(streamIds -> {
  					ArrayList<T> sortedList = input.list;
  					input.list = null;
  					return Promise.ofBlocking(sortingExecutor, () -> sortedList.sort(itemComparator))
  						.map($ -> {
  							StreamSupplier<T> listSupplier = StreamSuppliers.ofIterator(deduplicate ?
  								new DistinctIterator<>(sortedList, keyFunction, keyComparator) :
  								sortedList.iterator());
  							logger.info("Items in memory: {}, files: {}", sortedList.size(), streamIds.size());
  							if (streamIds.isEmpty()) {
  								return listSupplier;
  							}
  							StreamReducer<K, T, Void> streamMerger = StreamReducer.create(keyComparator);
  							listSupplier.streamTo(streamMerger.newInput(keyFunction, deduplicate ? deduplicateReducer() : mergeReducer()));
  							for (Integer streamId : streamIds) {
  								StreamSuppliers.ofPromise(storage.read(streamId))
  									.streamTo(streamMerger.newInput(keyFunction, deduplicate ? deduplicateReducer() : mergeReducer()));
  							}
  							return streamMerger.getOutput();
  						});
  				}));
  	}
  
  	/**
  	 * Creates a new instance of StreamSorter
  	 *
  	 * @param storage           storage for storing elements which were not placed
  	 *                          to RAM
  	 * @param keyFunction       function for searching key
  	 * @param keyComparator     comparator for comparing key
  	 * @param distinct          if it is true it means that in result will be
  	 *                          not objects with same key
  	 * @param itemsInMemorySize size of elements which can be saved in RAM
  	 *                          before sorting
  	 */
  	public static <K, T> StreamSorter<K, T> create(
  		IStreamSorterStorage<T> storage, Function<T, K> keyFunction, Comparator<K> keyComparator, boolean distinct,
  		int itemsInMemorySize
  	) {
  		return StreamSorter.builder(storage, keyFunction, keyComparator, distinct, itemsInMemorySize).build();
  	}
  
  	/**
  	 * Creates a builder of StreamSorter
  	 *
  	 * @param storage           storage for storing elements which were not placed
  	 *                          to RAM
  	 * @param keyFunction       function for searching key
  	 * @param keyComparator     comparator for comparing key
  	 * @param distinct          if it is true it means that in result will be
  	 *                          not objects with same key
  	 * @param itemsInMemorySize size of elements which can be saved in RAM
  	 *                          before sorting
  	 */
  	public static <K, T> StreamSorter<K, T>.Builder builder(
  		IStreamSorterStorage<T> storage, Function<T, K> keyFunction, Comparator<K> keyComparator, boolean distinct,
  		int itemsInMemorySize
  	) {
  		return new StreamSorter<>(storage, keyFunction, keyComparator, distinct, itemsInMemorySize).new Builder();
  	}
  
  	public final class Builder extends AbstractBuilder<Builder, StreamSorter<K, T>> {
  		private Builder() {}
  
  		public Builder withSortingExecutor(Executor executor) {
  			checkNotBuilt(this);
  			sortingExecutor = executor;
  			return this;
  		}
  
  		@Override
  		protected StreamSorter<K, T> doBuild() {
  			return StreamSorter.this;
  		}
  	}
  
  	public final class Input extends AbstractStreamConsumer<T> implements StreamDataAcceptor<T> {
  		private final List<Integer> partitionIds;
  
  		private ArrayList<T> list = new ArrayList<>();
  		private Promise<Void> cleanupPromise;
  
  		private Input(List<Integer> partitionIds) {
  			this.partitionIds = partitionIds;
  		}
  
  		@Override
  		protected void onStarted() {
  			output.getAcknowledgement()
  				.then((ackRes, e) -> cleanup()
  					.then(($, e1) -> Promise.of(ackRes, e)))
  				.whenResult(this::acknowledge)
  				.whenException(this::closeEx);
  
  			resume(this);
  		}
  
  		@Override
  		public void accept(T item) {
  			list.add(item);
  			if (list.size() < itemsInMemory) {
  				return;
  			}
  			ArrayList<T> sortedList = this.list;
  			list = new ArrayList<>(itemsInMemory);
  
  			temporaryStreamsAccumulator.addPromise(
  				Promise.ofBlocking(sortingExecutor, () -> sortedList.sort(itemComparator))
  					.then($ -> {
  						Iterator<T> iterator = distinct ?
  							new DistinctIterator<>(sortedList, keyFunction, keyComparator) :
  							sortedList.iterator();
  						return storage.newPartitionId()
  							.then(partitionId -> storage.write(partitionId)
  								.then(consumer -> StreamSuppliers.ofIterator(iterator).streamTo(consumer))
  								.map($2 -> partitionId));
  					})
  					.whenResult(this::suspendOrResume)
  					.whenException(this::closeEx), List::add);
  
  			suspendOrResume();
  		}
  
  		private void suspendOrResume() {
  			if (temporaryStreamsAccumulator.getActivePromises() > 2) {
  				suspend();
  			} else {
  				resume(this);
  			}
  		}
  
  		@Override
  		protected void onEndOfStream() {
  			temporaryStreamsAccumulator.run();
  		}
  
  		@Override
  		protected void onError(Exception e) {
  			temporaryStreamsAccumulator.closeEx(e);
  		}
  
  		@Override
  		protected void onCleanup() {
  			list = null;
  			cleanup();
  		}
  
  		private Promise<Void> cleanup() {
  			if (cleanupPromise != null) return cleanupPromise;
  			return cleanupPromise = partitionIds.isEmpty() ?
  				Promise.complete() :
  				storage.cleanup(partitionIds);
  		}
  	}
  
  	@Override
  	public StreamConsumer<T> getInput() {
  		return input;
  	}
  
  	@Override
  	public StreamSupplier<T> getOutput() {
  		return output;
  	}
  
  }
  