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.cube;
  
  import io.activej.common.Utils;
  import io.activej.cube.aggregation.AggregationChunk;
  import io.activej.cube.aggregation.PrimaryKey;
  import io.activej.cube.aggregation.RangeTree;
  import io.activej.cube.aggregation.RangeTree.Segment;
  import io.activej.cube.aggregation.ot.AggregationDiff;
  import io.activej.cube.aggregation.predicate.AggregationPredicate;
  import io.activej.jmx.api.attribute.JmxAttribute;
  import io.activej.ot.OTState;
  import org.jetbrains.annotations.Nullable;
  import org.jetbrains.annotations.VisibleForTesting;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.*;
  import java.util.stream.Collectors;
  import java.util.stream.Stream;
  
  import static io.activej.common.Checks.checkArgument;
  import static io.activej.common.collection.CollectionUtils.intersection;
  import static io.activej.cube.aggregation.predicate.AggregationPredicates.*;
  import static java.util.Collections.unmodifiableMap;
  
  /**
   * Represents aggregation metadata. Chunks are stored in an index (represented by an array of {@link RangeTree}) for efficient search.
   * Provides methods for managing index, querying for chunks by key, searching for chunks that are available for consolidation.
   */
  public final class AggregationState implements OTState<AggregationDiff> {
  	private static final Logger logger = LoggerFactory.getLogger(AggregationState.class);
  
  	private final AggregationStructure structure;
  	private final Map<Long, AggregationChunk> chunks = new LinkedHashMap<>();
  	private RangeTree<PrimaryKey, AggregationChunk>[] prefixRanges;
  
  	private static final Comparator<AggregationChunk> MIN_KEY_ASCENDING_COMPARATOR = Comparator.comparing(AggregationChunk::getMinPrimaryKey);
  
  	AggregationState(AggregationStructure structure) {
  		this.structure = structure;
  		initIndex();
  	}
  
  	public AggregationStructure getStructure() {
  		return structure;
  	}
  
  	public Map<Long, AggregationChunk> getChunks() {
  		return unmodifiableMap(chunks);
  	}
  
  	@Override
  	public void apply(AggregationDiff commit) {
  		checkArgument(intersection(commit.getAddedChunks(), commit.getRemovedChunks()), Set::isEmpty,
  			v ->
  				"Non-empty intersection between added and removed chunks: " + v +
  				"\n Added chunks " + Utils.toString(commit.getAddedChunks()) +
  				"\n Removed chunks intersection: " + Utils.toString(commit.getRemovedChunks()));
  
  		for (AggregationChunk chunk : commit.getAddedChunks()) {
  			addToIndex(chunk);
  		}
  
  		for (AggregationChunk chunk : commit.getRemovedChunks()) {
  			removeFromIndex(chunk);
  		}
  	}
  
  	public double estimateCost(List<String> measures, AggregationPredicate predicate) {
  		return findChunks(measures, predicate).size();
  	}
  
  	public void addToIndex(AggregationChunk chunk) {
  		checkArgument(chunks.put(chunk.getChunkId(), chunk) == null,
  			() ->
  				"Trying to add existing chunk: " + chunk +
  				"\n this: " + this +
  				"\n chunks: " + Utils.toString(chunks.keySet()));
  
  		for (int size = 0; size <= structure.getKeys().size(); size++) {
  			RangeTree<PrimaryKey, AggregationChunk> index = prefixRanges[size];
  
  			PrimaryKey lower = chunk.getMinPrimaryKey().prefix(size);
  			PrimaryKey upper = chunk.getMaxPrimaryKey().prefix(size);
  			index.put(lower, upper, chunk);
  		}
  	}
  
  	public void removeFromIndex(AggregationChunk chunk) {
  		AggregationChunk removedChunk = chunks.remove(chunk.getChunkId());
  		checkArgument(removedChunk != null,
  			() ->
  				"Trying to remove unknown chunk: " + chunk +
  				"\n this: " + this +
  				"\n chunks: " + Utils.toString(chunks.keySet()));
  
  		for (int size = 0; size <= structure.getKeys().size(); size++) {
  			RangeTree<PrimaryKey, AggregationChunk> index = prefixRanges[size];
  
  			PrimaryKey lower = removedChunk.getMinPrimaryKey().prefix(size);
  			PrimaryKey upper = removedChunk.getMaxPrimaryKey().prefix(size);
  			index.remove(lower, upper, removedChunk);
  		}
  	}
  
  	@SuppressWarnings("unchecked")
  	void initIndex() {
  		prefixRanges = new RangeTree[structure.getKeys().size() + 1];
  		for (int size = 0; size <= structure.getKeys().size(); size++) {
  			prefixRanges[size] = RangeTree.create();
  		}
  	}
  
  	@Override
  	public void init() {
  		initIndex();
  		chunks.clear();
  	}
  
  	private static int getNumberOfOverlaps(Segment<?> segment) {
  		return segment.getSet().size() + segment.getClosingSet().size();
  	}
  
  	public Set<AggregationChunk> findOverlappingChunks() {
  		int minOverlaps = 2;
  		Set<AggregationChunk> result = new HashSet<>();
  		RangeTree<PrimaryKey, AggregationChunk> tree = prefixRanges[structure.getKeys().size()];
  		for (Map.Entry<PrimaryKey, Segment<AggregationChunk>> segmentEntry : tree.getSegments().entrySet()) {
  			Segment<AggregationChunk> segment = segmentEntry.getValue();
  			int overlaps = getNumberOfOverlaps(segment);
  			if (overlaps >= minOverlaps) {
  				result.addAll(segment.getSet());
  				result.addAll(segment.getClosingSet());
  			}
  		}
  		return result;
  	}
  
  	public List<AggregationChunk> findChunksGroupWithMostOverlaps() {
  		return findChunksGroupWithMostOverlaps(prefixRanges[structure.getKeys().size()], Set.of());
  	}
  
  	private static List<AggregationChunk> findChunksGroupWithMostOverlaps(RangeTree<PrimaryKey, AggregationChunk> tree, Set<Long> lockedChunkIds) {
  		int maxOverlaps = 2;
  		List<AggregationChunk> result = new ArrayList<>();
  		for (Map.Entry<PrimaryKey, Segment<AggregationChunk>> segmentEntry : tree.getSegments().entrySet()) {
  			Segment<AggregationChunk> segment = segmentEntry.getValue();
  			int overlaps = getNumberOfOverlaps(segment);
  			if (overlaps >= maxOverlaps) {
  				List<AggregationChunk> filteredChunks = filterOutChunks(segment, lockedChunkIds);
  
  				if (filteredChunks.size() < maxOverlaps) continue;
  
  				maxOverlaps = overlaps;
  				result = filteredChunks;
  			}
  		}
  		return result;
  	}
  
  	private static PickedChunks findChunksWithMinKeyOrSizeFixStrategy(
  		SortedMap<PrimaryKey, RangeTree<PrimaryKey, AggregationChunk>> partitioningKeyToTree, int maxChunks,
  		int optimalChunkSize, Set<Long> lockedChunkIds
  	) {
  		int minChunks = 2;
  		for (Map.Entry<PrimaryKey, RangeTree<PrimaryKey, AggregationChunk>> entry : partitioningKeyToTree.entrySet()) {
  			ChunksAndStrategy chunksAndStrategy = findChunksWithMinKeyOrSizeFixStrategy(entry.getValue(), maxChunks, optimalChunkSize, lockedChunkIds);
  			if (chunksAndStrategy.chunks.size() >= minChunks)
  				return new PickedChunks(chunksAndStrategy.strategy, entry.getValue(), chunksAndStrategy.chunks);
  		}
  		return new PickedChunks(PickingStrategy.MIN_KEY, null, List.of());
  	}
  
  	private static ChunksAndStrategy findChunksWithMinKeyOrSizeFixStrategy(
  		RangeTree<PrimaryKey, AggregationChunk> tree, int maxChunks, int optimalChunkSize, Set<Long> lockedChunkIds
  	) {
  		int minOverlaps = 2;
  		List<AggregationChunk> result = new ArrayList<>();
  		SortedMap<PrimaryKey, Segment<AggregationChunk>> tailMap = null;
  		for (Map.Entry<PrimaryKey, Segment<AggregationChunk>> segmentEntry : tree.getSegments().entrySet()) {
  			Segment<AggregationChunk> segment = segmentEntry.getValue();
  			List<AggregationChunk> segmentChunks = filterOutChunks(segment, lockedChunkIds);
  			int overlaps = segmentChunks.size();
  
  			// "min key" strategy
  			if (overlaps >= minOverlaps) {
  				result.addAll(segmentChunks);
  				return new ChunksAndStrategy(PickingStrategy.MIN_KEY, result);
  			}
  
  			// "size fix" strategy
  			if (overlaps == 1 && segmentChunks.get(0).getCount() != optimalChunkSize) {
  				tailMap = tree.getSegments().tailMap(segmentEntry.getKey());
  				break;
  			}
  		}
  
  		if (tailMap == null)
  			return new ChunksAndStrategy(PickingStrategy.SIZE_FIX, List.of());
  
  		Set<AggregationChunk> chunks = new HashSet<>();
  		for (Map.Entry<PrimaryKey, Segment<AggregationChunk>> segmentEntry : tailMap.entrySet()) {
  			if (chunks.size() >= maxChunks)
  				break;
  
  			Segment<AggregationChunk> segment = segmentEntry.getValue();
  			chunks.addAll(filterOutChunks(segment, lockedChunkIds));
  		}
  		result.addAll(chunks);
  
  		if (result.size() == 1) {
  			if (result.get(0).getCount() > optimalChunkSize)
  				return new ChunksAndStrategy(PickingStrategy.SIZE_FIX, result);
  			else
  				return new ChunksAndStrategy(PickingStrategy.SIZE_FIX, List.of());
  		}
  
  		return new ChunksAndStrategy(PickingStrategy.SIZE_FIX, result);
  	}
  
  	public enum PickingStrategy {
  		PARTITIONING,
  		HOT_SEGMENT,
  		MIN_KEY,
  		SIZE_FIX
  	}
  
  	public record PickedChunks(
  		PickingStrategy strategy, @Nullable RangeTree<PrimaryKey, AggregationChunk> partitionTree,
  		List<AggregationChunk> chunks
  	) {
  	}
  
  	public record ChunksAndStrategy(PickingStrategy strategy, List<AggregationChunk> chunks) {}
  
  	@VisibleForTesting
  	@Nullable SortedMap<PrimaryKey, RangeTree<PrimaryKey, AggregationChunk>> groupByPartition(int partitioningKeyLength) {
  		SortedMap<PrimaryKey, RangeTree<PrimaryKey, AggregationChunk>> partitioningKeyToTree = new TreeMap<>();
  
  		Set<AggregationChunk> allChunks = prefixRanges[0].getAll();
  		for (AggregationChunk chunk : allChunks) {
  			PrimaryKey minKeyPrefix = chunk.getMinPrimaryKey().prefix(partitioningKeyLength);
  			PrimaryKey maxKeyPrefix = chunk.getMaxPrimaryKey().prefix(partitioningKeyLength);
  
  			if (!minKeyPrefix.equals(maxKeyPrefix))
  				return null; // not partitioned
  
  			partitioningKeyToTree.computeIfAbsent(minKeyPrefix, $ -> RangeTree.create())
  				.put(chunk.getMinPrimaryKey(), chunk.getMaxPrimaryKey(), chunk);
  
  		}
  
  		return partitioningKeyToTree;
  	}
  
  	private List<AggregationChunk> findChunksForPartitioning(int partitioningKeyLength, int maxChunks, Set<Long> lockedChunkIds) {
  		List<AggregationChunk> chunksForPartitioning = new ArrayList<>();
  		List<AggregationChunk> allChunks = new ArrayList<>(prefixRanges[0].getAll());
  		allChunks.sort(MIN_KEY_ASCENDING_COMPARATOR);
  
  		for (AggregationChunk chunk : allChunks) {
  			if (chunksForPartitioning.size() == maxChunks)
  				break;
  
  			if (lockedChunkIds.contains(chunk.getChunkId()))
  				continue;
  
  			PrimaryKey minKeyPrefix = chunk.getMinPrimaryKey().prefix(partitioningKeyLength);
  			PrimaryKey maxKeyPrefix = chunk.getMaxPrimaryKey().prefix(partitioningKeyLength);
  
  			if (!minKeyPrefix.equals(maxKeyPrefix))
  				chunksForPartitioning.add(chunk);
  		}
  
  		return chunksForPartitioning;
  	}
  
  	public List<AggregationChunk> findChunksForConsolidationMinKey(int maxChunks, int optimalChunkSize) {
  		return findChunksForConsolidationMinKey(maxChunks, optimalChunkSize, Set.of());
  	}
  
  	public List<AggregationChunk> findChunksForConsolidationMinKey(int maxChunks, int optimalChunkSize, Set<Long> lockedChunkIds) {
  		int partitioningKeyLength = structure.getPartitioningKey().size();
  		SortedMap<PrimaryKey, RangeTree<PrimaryKey, AggregationChunk>> partitioningKeyToTree = groupByPartition(partitioningKeyLength);
  		if (partitioningKeyToTree == null) { // not partitioned
  			List<AggregationChunk> chunks = findChunksForPartitioning(partitioningKeyLength, maxChunks, lockedChunkIds);
  			logChunksAndStrategy(chunks, PickingStrategy.PARTITIONING);
  			return chunks; // launch partitioning
  		}
  		PickedChunks pickedChunks = findChunksWithMinKeyOrSizeFixStrategy(partitioningKeyToTree, maxChunks, optimalChunkSize, lockedChunkIds);
  		return processSelection(pickedChunks.chunks, maxChunks, pickedChunks.partitionTree, pickedChunks.strategy, lockedChunkIds);
  	}
  
  	public List<AggregationChunk> findChunksForConsolidationHotSegment(int maxChunks) {
  		return findChunksForConsolidationHotSegment(maxChunks, Set.of());
  	}
  
  	public List<AggregationChunk> findChunksForConsolidationHotSegment(int maxChunks, Set<Long> lockedChunkIds) {
  		RangeTree<PrimaryKey, AggregationChunk> tree = prefixRanges[prefixRanges.length - 1];
  		List<AggregationChunk> chunks = findChunksGroupWithMostOverlaps(tree, lockedChunkIds);
  		return processSelection(chunks, maxChunks, tree, PickingStrategy.HOT_SEGMENT, lockedChunkIds);
  	}
  
  	private static List<AggregationChunk> filterOutChunks(Segment<AggregationChunk> segment, Set<Long> lockedChunkIds) {
  		return Stream.concat(segment.getSet().stream(), segment.getClosingSet().stream())
  			.filter(aggregationChunk -> !lockedChunkIds.contains(aggregationChunk.getChunkId()))
  			.collect(Collectors.toList());
  	}
  
  	private static List<AggregationChunk> processSelection(
  		List<AggregationChunk> chunks, int maxChunks, RangeTree<PrimaryKey, AggregationChunk> partitionTree,
  		PickingStrategy strategy, Set<Long> lockedChunkIds
  	) {
  		if (chunks.isEmpty() || chunks.size() == maxChunks) {
  			logChunksAndStrategy(chunks, strategy);
  			return chunks;
  		}
  
  		if (chunks.size() > maxChunks) {
  			List<AggregationChunk> trimmedChunks = trimChunks(chunks, maxChunks);
  			logChunksAndStrategy(trimmedChunks, strategy);
  			return trimmedChunks;
  		}
  
  		if (strategy == PickingStrategy.SIZE_FIX) {
  			logChunksAndStrategy(chunks, PickingStrategy.SIZE_FIX);
  			return chunks;
  		}
  
  		List<AggregationChunk> expandedChunks = expandRange(partitionTree, chunks, maxChunks, lockedChunkIds);
  
  		if (expandedChunks.size() > maxChunks) {
  			List<AggregationChunk> trimmedChunks = trimChunks(expandedChunks, maxChunks);
  			logChunksAndStrategy(trimmedChunks, strategy);
  			return trimmedChunks;
  		}
  
  		logChunksAndStrategy(expandedChunks, strategy);
  		return expandedChunks;
  	}
  
  	private static void logChunksAndStrategy(Collection<AggregationChunk> chunks, PickingStrategy strategy) {
  		if (logger.isInfoEnabled()) {
  			String chunkIds = chunks.stream()
  				.map(AggregationChunk::getChunkId)
  				.map(String::valueOf)
  				.collect(Collectors.joining(",", "[", "]"));
  			logger.info("Chunks for consolidation {}: {}. Strategy: {}", chunks.size(), chunkIds, strategy);
  		}
  	}
  
  	private static List<AggregationChunk> trimChunks(List<AggregationChunk> chunks, int maxChunks) {
  		chunks.sort(MIN_KEY_ASCENDING_COMPARATOR);
  		return chunks.subList(0, maxChunks);
  	}
  
  	private static boolean expandRange(RangeTree<PrimaryKey, AggregationChunk> tree, Set<AggregationChunk> chunks, Set<Long> lockedChunkIds) {
  		PrimaryKey minKey = null;
  		PrimaryKey maxKey = null;
  
  		for (AggregationChunk chunk : chunks) {
  			PrimaryKey chunkMinKey = chunk.getMinPrimaryKey();
  			PrimaryKey chunkMaxKey = chunk.getMaxPrimaryKey();
  
  			if (minKey == null) {
  				minKey = chunkMinKey;
  				maxKey = chunkMaxKey;
  				continue;
  			}
  
  			if (chunkMinKey.compareTo(minKey) < 0)
  				minKey = chunkMinKey;
  
  			if (chunkMaxKey.compareTo(maxKey) > 0)
  				maxKey = chunkMaxKey;
  		}
  
  		boolean expanded = false;
  		for (AggregationChunk aggregationChunk : tree.getRange(minKey, maxKey)) {
  			if (!lockedChunkIds.contains(aggregationChunk.getChunkId()) && chunks.add(aggregationChunk)) {
  				expanded = true;
  			}
  		}
  		return expanded;
  	}
  
  	private static void expandRange(RangeTree<PrimaryKey, AggregationChunk> tree, Set<AggregationChunk> chunks, int maxChunks, Set<Long> lockedChunkIds) {
  		boolean expand = chunks.size() < maxChunks;
  
  		while (expand) {
  			boolean expanded = expandRange(tree, chunks, lockedChunkIds);
  			expand = expanded && chunks.size() < maxChunks;
  		}
  	}
  
  	private static List<AggregationChunk> expandRange(
  		RangeTree<PrimaryKey, AggregationChunk> tree, List<AggregationChunk> chunks, int maxChunks,
  		Set<Long> lockedChunkIds
  	) {
  		Set<AggregationChunk> chunkSet = new HashSet<>(chunks);
  		expandRange(tree, chunkSet, maxChunks, lockedChunkIds);
  		return new ArrayList<>(chunkSet);
  	}
  
  	public List<ConsolidationDebugInfo> getConsolidationDebugInfo() {
  		List<ConsolidationDebugInfo> infos = new ArrayList<>();
  		RangeTree<PrimaryKey, AggregationChunk> tree = prefixRanges[prefixRanges.length - 1];
  
  		for (Map.Entry<PrimaryKey, Segment<AggregationChunk>> segmentEntry : tree.getSegments().entrySet()) {
  			PrimaryKey key = segmentEntry.getKey();
  			Segment<AggregationChunk> segment = segmentEntry.getValue();
  			int overlaps = segment.getSet().size() + segment.getClosingSet().size();
  			Set<AggregationChunk> segmentSet = segment.getSet();
  			Set<AggregationChunk> segmentClosingSet = segment.getClosingSet();
  			infos.add(new ConsolidationDebugInfo(key, segmentSet, segmentClosingSet, overlaps));
  		}
  
  		return infos;
  	}
  
  	public record ConsolidationDebugInfo(
  			PrimaryKey key, Set<AggregationChunk> segmentSet, Set<AggregationChunk> segmentClosingSet, int overlaps
  	) {
  	}
  
  	@VisibleForTesting
  	public static boolean chunkMightContainQueryValues(
  		PrimaryKey minQueryKey, PrimaryKey maxQueryKey, PrimaryKey minChunkKey, PrimaryKey maxChunkKey
  	) {
  		return chunkMightContainQueryValues(minQueryKey.values(), maxQueryKey.values(),
  			minChunkKey.values(), maxChunkKey.values());
  	}
  
  	@SuppressWarnings("unchecked")
  	private static boolean chunkMightContainQueryValues(
  		List<Object> queryMinValues, List<Object> queryMaxValues, List<Object> chunkMinValues, List<Object> chunkMaxValues
  	) {
  		checkArgument(queryMinValues.size() == queryMaxValues.size(),
  			"Sizes of lists of query minimum and maximum values should match");
  		checkArgument(chunkMinValues.size() == chunkMaxValues.size(),
  			"Sizes of lists of chunk minimum and maximum values should match");
  
  		for (int i = 0; i < queryMinValues.size(); ++i) {
  			Comparable<Object> queryMinValue = (Comparable<Object>) queryMinValues.get(i);
  			Comparable<Object> queryMaxValue = (Comparable<Object>) queryMaxValues.get(i);
  			Comparable<Object> chunkMinValue = (Comparable<Object>) chunkMinValues.get(i);
  			Comparable<Object> chunkMaxValue = (Comparable<Object>) chunkMaxValues.get(i);
  
  			if (chunkMinValue.compareTo(chunkMaxValue) == 0) {
  				if (!(queryMinValue.compareTo(chunkMinValue) <= 0 && queryMaxValue.compareTo(chunkMaxValue) >= 0)) {
  					return false;
  				}
  			} else {
  				return queryMinValue.compareTo(chunkMaxValue) <= 0 && queryMaxValue.compareTo(chunkMinValue) >= 0;
  			}
  		}
  
  		return true;
  	}
  
  	public List<AggregationChunk> findChunks(List<String> measures, AggregationPredicate predicate) {
  		RangeScan rangeScan = toRangeScan(predicate, structure.getKeys(), structure.getKeyTypes());
  		if (rangeScan.isNoScan())
  			return List.of();
  
  		Set<String> requestedFields = new HashSet<>(structure.findFields(measures));
  		List<AggregationChunk> chunks = new ArrayList<>();
  		for (AggregationChunk chunk : rangeQuery(rangeScan.getFrom(), rangeScan.getTo())) {
  			if (intersection(new HashSet<>(chunk.getMeasures()), requestedFields).isEmpty())
  				continue;
  
  			AggregationPredicate chunkPredicate = chunk.toPredicate(structure.getKeys(), structure.getKeyTypes());
  			if (and(chunkPredicate, predicate).simplify().equals(alwaysFalse())) {
  				continue;
  			}
  
  			chunks.add(chunk);
  		}
  
  		return chunks;
  	}
  
  	private List<AggregationChunk> rangeQuery(PrimaryKey minPrimaryKey, PrimaryKey maxPrimaryKey) {
  		checkArgument(minPrimaryKey.size() == maxPrimaryKey.size(), "Sizes of min primary key and max primary key should match");
  		int size = minPrimaryKey.size();
  		RangeTree<PrimaryKey, AggregationChunk> index = prefixRanges[size];
  		return new ArrayList<>(index.getRange(minPrimaryKey, maxPrimaryKey));
  	}
  
  	@JmxAttribute
  	public int getNumberOfOverlappingChunks() {
  		return findOverlappingChunks().size();
  	}
  
  	@JmxAttribute(name = "chunks")
  	public int getChunksSize() {
  		return chunks.size();
  	}
  
  	@Override
  	public String toString() {
  		return "AggregationState{chunks = " + getChunksSize() + "}";
  	}
  }
  