/*
 * 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.crdt.storage.cluster;

import io.activej.common.ApplicationSettings;
import io.activej.common.HashUtils;
import org.jetbrains.annotations.Nullable;

import java.util.*;
import java.util.function.ToIntFunction;

import static io.activej.common.Checks.checkArgument;

public final class RendezvousHashSharder<K> implements Sharder<K> {
	public static final int NUMBER_OF_BUCKETS = ApplicationSettings.getInt(RendezvousHashSharder.class, "numberOfBuckets", 512);

	static {
		checkArgument((NUMBER_OF_BUCKETS & (NUMBER_OF_BUCKETS - 1)) == 0, "Number of buckets must be a power of two");
	}

	final int[][] buckets;
	final ToIntFunction<K> keyHashFn;

	private RendezvousHashSharder(int[][] buckets, ToIntFunction<K> keyHashFn) {
		this.buckets = buckets;
		this.keyHashFn = keyHashFn;
	}

	public static <K extends Comparable<K>, P> RendezvousHashSharder<K> create(
		ToIntFunction<K> keyHashFn, ToIntFunction<P> partitionIdHashCode, Set<P> partitions, List<P> partitionsAlive,
		int shards, boolean repartition
	) {
		Map<P, Integer> partitionsAliveMap = new HashMap<>();
		for (P partitionId : partitionsAlive) {
			partitionsAliveMap.put(partitionId, partitionsAliveMap.size());
		}

		int[][] buckets = new int[NUMBER_OF_BUCKETS][];

		//noinspection NullableProblems
		class ObjWithIndex {
			final P partitionId;
			final @Nullable Integer aliveIndex;
			private long hash;

			public long getHash() {
				return hash;
			}

			ObjWithIndex(P partitionId, Integer aliveIndex) {
				this.partitionId = partitionId;
				this.aliveIndex = aliveIndex;
			}
		}

		ObjWithIndex[] toSort = new ObjWithIndex[partitions.size()];
		int i = 0;
		for (P partitionId : partitions) {
			toSort[i++] = new ObjWithIndex(partitionId, partitionsAliveMap.get(partitionId));
		}
		int[] buf = new int[partitions.size()];

		for (int bucket = 0; bucket < buckets.length; bucket++) {
			for (ObjWithIndex obj : toSort) {
				obj.hash = hashBucket(partitionIdHashCode.applyAsInt(obj.partitionId), bucket);
			}

			Arrays.sort(toSort, Comparator.comparingLong(ObjWithIndex::getHash).reversed());

			int n = 0;
			int shardsN = shards;
			for (ObjWithIndex obj : toSort) {
				if (shardsN <= 0) break;
				if (obj.aliveIndex != null) {
					buf[n++] = obj.aliveIndex;
					shardsN--;
				} else if (!repartition) {
					shardsN--;
				}
			}

			buckets[bucket] = Arrays.copyOf(buf, n);
			Arrays.sort(buckets[bucket]);
		}
		return new RendezvousHashSharder<>(buckets, keyHashFn);
	}

	static <K> Sharder<K> unionOf(List<RendezvousHashSharder<K>> sharders) {
		if (sharders.isEmpty()) return Sharder.none();
		if (sharders.size() == 1) return sharders.get(0);
		int[][] buckets = new int[sharders.get(0).buckets.length][];
		ToIntFunction<K> keyHashFn = sharders.get(0).keyHashFn;
		int[] buf = new int[0];
		for (int bucket = 0; bucket < buckets.length; bucket++) {
			int pos = 0;
			for (RendezvousHashSharder<K> sharder : sharders) {
				int[] selected = sharder.buckets[bucket];
				NEXT:
				for (int idx : selected) {
					for (int i = 0; i < pos; i++) {
						if (idx == buf[i]) continue NEXT;
					}
					if (buf.length <= pos) buf = Arrays.copyOf(buf, buf.length * 2 + 1);
					buf[pos++] = idx;
				}
			}
			buckets[bucket] = Arrays.copyOf(buf, pos);
		}
		return new RendezvousHashSharder<>(buckets, keyHashFn);
	}

	@Override
	public int[] shard(K key) {
		return buckets[keyHashFn.applyAsInt(key) & (NUMBER_OF_BUCKETS - 1)];
	}

	public static long hashBucket(int partitionHash, int bucket) {
		return HashUtils.murmur3hash(((long) partitionHash << 32) | (bucket & 0xFFFFFFFFL));
	}
}
