/*
 * 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.cassandra.index.sai.memory;

import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.SortedSet;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Function;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.memtable.TrieMemtable;
import org.apache.cassandra.db.tries.InMemoryTrie;
import org.apache.cassandra.db.tries.Trie;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.QueryContext;
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.analyzer.AbstractAnalyzer;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentMetadata;
import org.apache.cassandra.index.sai.disk.v1.vector.PrimaryKeyWithScore;
import org.apache.cassandra.index.sai.iterators.KeyRangeIterator;
import org.apache.cassandra.index.sai.plan.Expression;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.index.sai.utils.PrimaryKeys;
import org.apache.cassandra.utils.CloseableIterator;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;

/**
 * This is an in-memory index using the {@link InMemoryTrie} to store a {@link ByteComparable}
 * representation of the indexed values. Data is stored on-heap or off-heap and follows the
 * settings of the {@link TrieMemtable} to determine where.
 */
public class TrieMemoryIndex extends MemoryIndex
{
    private static final Logger logger = LoggerFactory.getLogger(TrieMemoryIndex.class);
    private static final int MAX_RECURSIVE_KEY_LENGTH = 128;
    private static final int MINIMUM_PRIORITY_QUEUE_SIZE = 128;

    private final InMemoryTrie<PrimaryKeys> data;
    private final PrimaryKeysReducer primaryKeysReducer;

    private ByteBuffer minTerm;
    private ByteBuffer maxTerm;

    // Maintain the last queue size used on this index to use for the next range match.
    // This allows for receiving a stream of wide range queries where the queue size
    // is larger than we would want to default the size to.
    private final AtomicInteger lastPriorityQueueSize = new AtomicInteger(MINIMUM_PRIORITY_QUEUE_SIZE);

    public TrieMemoryIndex(StorageAttachedIndex index)
    {
        super(index);
        this.data = new InMemoryTrie<>(TrieMemtable.BUFFER_TYPE);
        this.primaryKeysReducer = new PrimaryKeysReducer();
    }

    /**
     * Adds an index value to the in-memory index
     *
     * @param key partition key for the indexed value
     * @param clustering clustering for the indexed value
     * @param value indexed value
     * @return amount of heap allocated by the new value
     */
    @Override
    public synchronized long add(DecoratedKey key, Clustering<?> clustering, ByteBuffer value)
    {
        value = index.termType().asIndexBytes(value);
        final PrimaryKey primaryKey = index.hasClustering() ? index.keyFactory().create(key, clustering)
                                                            : index.keyFactory().create(key);
        final long initialSizeOnHeap = data.sizeOnHeap();
        final long initialSizeOffHeap = data.sizeOffHeap();
        final long reducerHeapSize = primaryKeysReducer.heapAllocations();

        if (index.hasAnalyzer())
        {
            AbstractAnalyzer analyzer = index.analyzer();
            try
            {
                analyzer.reset(value);
                while (analyzer.hasNext())
                {
                    addTerm(primaryKey, analyzer.next());
                }
            }
            finally
            {
                analyzer.end();
            }
        }
        else
        {
            addTerm(primaryKey, value);
        }
        long onHeap = data.sizeOnHeap();
        long offHeap = data.sizeOffHeap();
        long heapAllocations = primaryKeysReducer.heapAllocations();
        return (onHeap - initialSizeOnHeap) + (offHeap - initialSizeOffHeap) + (heapAllocations - reducerHeapSize);
    }

    @Override
    public long update(DecoratedKey key, Clustering<?> clustering, ByteBuffer oldValue, ByteBuffer newValue)
    {
        throw new UnsupportedOperationException();
    }

    /**
     * Search for an expression in the in-memory index within the {@link AbstractBounds} defined
     * by keyRange. This can either be an exact match or a range match.
     * <p>
     * @param expression the {@link Expression} to search for
     * @param keyRange the {@link AbstractBounds} containing the key range to restrict the search to
     * @return a {@link KeyRangeIterator} containing the search results
     */
    public KeyRangeIterator search(QueryContext queryContext, Expression expression, AbstractBounds<PartitionPosition> keyRange)
    {
        logger.trace("Searching memtable index on expression '{}'...", expression);

        switch (expression.getIndexOperator())
        {
            case EQ:
            case CONTAINS_KEY:
            case CONTAINS_VALUE:
                return exactMatch(expression, keyRange);
            case RANGE:
                KeyRangeIterator keyIterator = rangeMatch(expression, keyRange);
                int keyCount = (int) keyIterator.getMaxKeys();
                if (keyCount > MINIMUM_PRIORITY_QUEUE_SIZE)
                    lastPriorityQueueSize.set(keyCount);
                return keyIterator;
            default:
                throw new IllegalArgumentException("Unsupported expression: " + expression);
        }
    }

    /**
     * Returns an {@link Iterator} over the entire dataset contained in the trie. This is used
     * when the index is flushed to disk.
     *
     * @return the iterator containing the trie data
     */
    @Override
    public Iterator<Pair<ByteComparable, PrimaryKeys>> iterator()
    {
        Iterator<Map.Entry<ByteComparable, PrimaryKeys>> iterator = data.entrySet().iterator();
        return new Iterator<>()
        {
            @Override
            public boolean hasNext()
            {
                return iterator.hasNext();
            }

            @Override
            public Pair<ByteComparable, PrimaryKeys> next()
            {
                Map.Entry<ByteComparable, PrimaryKeys> entry = iterator.next();
                return Pair.create(entry.getKey(), entry.getValue());
            }
        };
    }

    @Override
    public SegmentMetadata.ComponentMetadataMap writeDirect(IndexDescriptor indexDescriptor,
                                                            IndexIdentifier indexIdentifier,
                                                            Function<PrimaryKey, Integer> postingTransformer)
    {
        throw new UnsupportedOperationException();
    }

    @Override
    public boolean isEmpty()
    {
        return minTerm == null;
    }

    @Override
    public ByteBuffer getMinTerm()
    {
        return minTerm;
    }

    @Override
    public ByteBuffer getMaxTerm()
    {
        return maxTerm;
    }

    private void addTerm(PrimaryKey primaryKey, ByteBuffer term)
    {
        if (index.validateTermSize(primaryKey.partitionKey(), term, false, null))
        {
            setMinMaxTerm(term.duplicate());

            final ByteComparable comparableBytes = asComparableBytes(term);

            try
            {
                if (term.limit() <= MAX_RECURSIVE_KEY_LENGTH)
                {
                    data.putRecursive(comparableBytes, primaryKey, primaryKeysReducer);
                }
                else
                {
                    data.apply(Trie.singleton(comparableBytes, primaryKey), primaryKeysReducer);
                }
            }
            catch (InMemoryTrie.SpaceExhaustedException e)
            {
                throw new RuntimeException(e);
            }
        }
    }

    private void setMinMaxTerm(ByteBuffer term)
    {
        assert term != null;

        minTerm = index.termType().min(term, minTerm);
        maxTerm = index.termType().max(term, maxTerm);
    }

    private ByteComparable asComparableBytes(ByteBuffer input)
    {
        return version -> index.termType().asComparableBytes(input, version);
    }

    private KeyRangeIterator exactMatch(Expression expression, AbstractBounds<PartitionPosition> keyRange)
    {
        ByteComparable comparableMatch = expression.lower() == null ? ByteComparable.EMPTY
                                                                    : asComparableBytes(expression.lower().value.encoded);
        PrimaryKeys primaryKeys = data.get(comparableMatch);
        return primaryKeys == null ? KeyRangeIterator.empty()
                                   : new FilteringInMemoryKeyRangeIterator(primaryKeys.keys(), keyRange);
    }

    @Override
    public CloseableIterator<PrimaryKeyWithScore> orderBy(QueryContext queryContext, Expression orderer, AbstractBounds<PartitionPosition> keyRange)
    {
        throw new UnsupportedOperationException();
    }

    @Override
    public CloseableIterator<PrimaryKeyWithScore> orderResultsBy(QueryContext context, List<PrimaryKey> results, Expression orderer)
    {
        throw new UnsupportedOperationException();
    }

    private static class Collector
    {
        final PriorityQueue<PrimaryKey> mergedKeys;
        final AbstractBounds<PartitionPosition> keyRange;

        PrimaryKey maximumKey = null;

        public Collector(AbstractBounds<PartitionPosition> keyRange, int expectedKeys)
        {
            this.keyRange = keyRange;
            this.mergedKeys = new PriorityQueue<>(expectedKeys);
        }

        public void processContent(PrimaryKeys keys)
        {
            if (keys.isEmpty())
                return;

            SortedSet<PrimaryKey> primaryKeys = keys.keys();

            // shortcut to avoid generating iterator
            if (primaryKeys.size() == 1)
            {
                processKey(primaryKeys.first());
                return;
            }

            // skip entire partition keys if they don't overlap
            if (!keyRange.right.isMinimum() && primaryKeys.first().partitionKey().compareTo(keyRange.right) > 0
                || primaryKeys.last().partitionKey().compareTo(keyRange.left) < 0)
                return;

            for (PrimaryKey primaryKey : primaryKeys)
                processKey(primaryKey);
        }

        private void processKey(PrimaryKey key)
        {
            if (keyRange.contains(key.partitionKey()))
            {
                mergedKeys.add(key);

                // We only track the maximum key, as the minimum can be peeked in constant time on the PQ itself.
                maximumKey = maximumKey == null ? key : key.compareTo(maximumKey) > 0 ? key : maximumKey;
            }
        }
    }

    private KeyRangeIterator rangeMatch(Expression expression, AbstractBounds<PartitionPosition> keyRange)
    {
        ByteComparable lowerBound, upperBound;
        boolean lowerInclusive, upperInclusive;
        if (expression.lower() != null)
        {
            lowerBound = asComparableBytes(expression.lower().value.encoded);
            lowerInclusive = expression.lower().inclusive;
        }
        else
        {
            lowerBound = ByteComparable.EMPTY;
            lowerInclusive = false;
        }

        if (expression.upper() != null)
        {
            upperBound = asComparableBytes(expression.upper().value.encoded);
            upperInclusive = expression.upper().inclusive;
        }
        else
        {
            upperBound = null;
            upperInclusive = false;
        }

        Collector cd = new Collector(keyRange, lastPriorityQueueSize.get());
        Iterator<PrimaryKeys> values = data.subtrie(lowerBound, lowerInclusive, upperBound, upperInclusive).valueIterator();

        while (values.hasNext())
            cd.processContent(values.next());

        if (cd.mergedKeys.isEmpty())
            return KeyRangeIterator.empty();

        return new InMemoryKeyRangeIterator(cd.mergedKeys.peek(), cd.maximumKey, cd.mergedKeys);
    }

    private static class PrimaryKeysReducer implements InMemoryTrie.UpsertTransformer<PrimaryKeys, PrimaryKey>
    {
        private final LongAdder heapAllocations = new LongAdder();

        @Override
        public PrimaryKeys apply(PrimaryKeys existing, PrimaryKey neww)
        {
            if (existing == null)
            {
                existing = new PrimaryKeys();
                heapAllocations.add(existing.unsharedHeapSize());
            }
            heapAllocations.add(existing.add(neww));
            return existing;
        }

        long heapAllocations()
        {
            return heapAllocations.longValue();
        }
    }
}
