/*
 * 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.flink.contrib.streaming.state.iterator;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.contrib.streaming.state.RocksIteratorWrapper;
import org.apache.flink.util.IOUtils;
import org.apache.flink.util.Preconditions;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

/**
 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
 * The resulting iteration sequence is ordered by (key-group, kv-state).
 */
public class RocksStatesPerKeyGroupMergeIterator implements AutoCloseable {

	private final PriorityQueue<RocksSingleStateIterator> heap;
	private final int keyGroupPrefixByteCount;
	private boolean newKeyGroup;
	private boolean newKVState;
	private boolean valid;
	private RocksSingleStateIterator currentSubIterator;

	private static final List<Comparator<RocksSingleStateIterator>> COMPARATORS;

	static {
		int maxBytes = 2;
		COMPARATORS = new ArrayList<>(maxBytes);
		for (int i = 0; i < maxBytes; ++i) {
			final int currentBytes = i + 1;
			COMPARATORS.add((o1, o2) -> {
				int arrayCmpRes = compareKeyGroupsForByteArrays(
					o1.getCurrentKey(), o2.getCurrentKey(), currentBytes);
				return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
			});
		}
	}

	public RocksStatesPerKeyGroupMergeIterator(
		List<Tuple2<RocksIteratorWrapper, Integer>> kvStateIterators,
		final int keyGroupPrefixByteCount) {
		Preconditions.checkNotNull(kvStateIterators);
		Preconditions.checkArgument(keyGroupPrefixByteCount >= 1);

		this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;

		if (kvStateIterators.size() > 0) {
			this.heap = buildIteratorHeap(kvStateIterators);
			this.valid = !heap.isEmpty();
			this.currentSubIterator = heap.poll();
			kvStateIterators.clear();
		} else {
			// creating a PriorityQueue of size 0 results in an exception.
			this.heap = null;
			this.valid = false;
		}

		this.newKeyGroup = true;
		this.newKVState = true;
	}

	/**
	 * Advances the iterator. Should only be called if {@link #isValid()} returned true.
	 * Valid flag can only change after calling {@link #next()}.
	 */
	public void next() {
		newKeyGroup = false;
		newKVState = false;

		final RocksIteratorWrapper rocksIterator = currentSubIterator.getIterator();
		rocksIterator.next();

		byte[] oldKey = currentSubIterator.getCurrentKey();
		if (rocksIterator.isValid()) {

			currentSubIterator.setCurrentKey(rocksIterator.key());

			if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
				heap.offer(currentSubIterator);
				currentSubIterator = heap.remove();
				newKVState = currentSubIterator.getIterator() != rocksIterator;
				detectNewKeyGroup(oldKey);
			}
		} else {
			IOUtils.closeQuietly(rocksIterator);

			if (heap.isEmpty()) {
				currentSubIterator = null;
				valid = false;
			} else {
				currentSubIterator = heap.remove();
				newKVState = true;
				detectNewKeyGroup(oldKey);
			}
		}
	}

	private PriorityQueue<RocksSingleStateIterator> buildIteratorHeap(
		List<Tuple2<RocksIteratorWrapper, Integer>> kvStateIterators) {

		Comparator<RocksSingleStateIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount - 1);

		PriorityQueue<RocksSingleStateIterator> iteratorPriorityQueue =
			new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);

		for (Tuple2<RocksIteratorWrapper, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
			final RocksIteratorWrapper rocksIterator = rocksIteratorWithKVStateId.f0;
			rocksIterator.seekToFirst();
			if (rocksIterator.isValid()) {
				iteratorPriorityQueue.offer(
					new RocksSingleStateIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
			} else {
				IOUtils.closeQuietly(rocksIterator);
			}
		}
		return iteratorPriorityQueue;
	}

	private boolean isDifferentKeyGroup(byte[] a, byte[] b) {
		return 0 != compareKeyGroupsForByteArrays(a, b, keyGroupPrefixByteCount);
	}

	private void detectNewKeyGroup(byte[] oldKey) {
		if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
			newKeyGroup = true;
		}
	}

	/**
	 * @return key-group for the current key
	 */
	public int keyGroup() {
		final byte[] currentKey = currentSubIterator.getCurrentKey();
		int result = 0;
		//big endian decode
		for (int i = 0; i < keyGroupPrefixByteCount; ++i) {
			result <<= 8;
			result |= (currentKey[i] & 0xFF);
		}
		return result;
	}

	public byte[] key() {
		return currentSubIterator.getCurrentKey();
	}

	public byte[] value() {
		return currentSubIterator.getIterator().value();
	}

	/**
	 * @return Id of K/V state to which the current key belongs.
	 */
	public int kvStateId() {
		return currentSubIterator.getKvStateId();
	}

	/**
	 * Indicates if current key starts a new k/v-state, i.e. belong to a different k/v-state than it's predecessor.
	 * @return true iff the current key belong to a different k/v-state than it's predecessor.
	 */
	public boolean isNewKeyValueState() {
		return newKVState;
	}

	/**
	 * Indicates if current key starts a new key-group, i.e. belong to a different key-group than it's predecessor.
	 * @return true iff the current key belong to a different key-group than it's predecessor.
	 */
	public boolean isNewKeyGroup() {
		return newKeyGroup;
	}

	/**
	 * Check if the iterator is still valid. Getters like {@link #key()}, {@link #value()}, etc. as well as
	 * {@link #next()} should only be called if valid returned true. Should be checked after each call to
	 * {@link #next()} before accessing iterator state.
	 * @return True iff this iterator is valid.
	 */
	public boolean isValid() {
		return valid;
	}

	private static int compareKeyGroupsForByteArrays(byte[] a, byte[] b, int len) {
		for (int i = 0; i < len; ++i) {
			int diff = (a[i] & 0xFF) - (b[i] & 0xFF);
			if (diff != 0) {
				return diff;
			}
		}
		return 0;
	}

	@Override
	public void close() {
		IOUtils.closeQuietly(currentSubIterator);
		currentSubIterator = null;

		IOUtils.closeAllQuietly(heap);
		heap.clear();
	}
}