/*
 * 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.disk.v1.segment;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicInteger;
import javax.annotation.concurrent.NotThreadSafe;

import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
import org.apache.cassandra.index.sai.disk.v1.bbtree.NumericIndexWriter;
import org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter;
import org.apache.cassandra.index.sai.disk.v1.vector.OnHeapGraph;
import org.apache.cassandra.index.sai.utils.NamedMemoryLimiter;
import org.apache.cassandra.index.sai.utils.PrimaryKey;

/**
 * Creates an on-heap index data structure to be flushed to an SSTable index.
 */
@NotThreadSafe
public abstract class SegmentBuilder
{
    private static final Logger logger = LoggerFactory.getLogger(SegmentBuilder.class);

    // Served as safe net in case memory limit is not triggered or when merger merges small segments..
    public static final long LAST_VALID_SEGMENT_ROW_ID = (Integer.MAX_VALUE / 2) - 1L;
    private static long testLastValidSegmentRowId = -1;

    /** The number of column indexes being built globally. */
    private static final AtomicInteger ACTIVE_BUILDER_COUNT = new AtomicInteger(0);

    /** Minimum flush size, dynamically updated as segment builds are started and completed/aborted. */
    private static volatile long minimumFlushBytes;
    private final NamedMemoryLimiter limiter;
    private final long lastValidSegmentRowID;
    private boolean flushed = false;
    private boolean active = true;
    // segment metadata
    private long minSSTableRowId = -1;
    private long maxSSTableRowId = -1;
    private long segmentRowIdOffset = 0;

    // in token order
    private PrimaryKey minKey;
    private PrimaryKey maxKey;
    // in termComparator order
    private ByteBuffer minTerm;
    private ByteBuffer maxTerm;

    final StorageAttachedIndex index;
    long totalBytesAllocated;
    int rowCount = 0;
    int maxSegmentRowId = -1;

    public static class TrieSegmentBuilder extends SegmentBuilder
    {
        protected final SegmentTrieBuffer segmentTrieBuffer;

        public TrieSegmentBuilder(StorageAttachedIndex index, NamedMemoryLimiter limiter)
        {
            super(index, limiter);

            segmentTrieBuffer = new SegmentTrieBuffer();
            totalBytesAllocated = segmentTrieBuffer.memoryUsed();
        }

        @Override
        protected long addInternal(ByteBuffer term, int segmentRowId)
        {
            return segmentTrieBuffer.add(v -> index.termType().asComparableBytes(term, v), term.limit(), segmentRowId);
        }

        @Override
        protected SegmentMetadata.ComponentMetadataMap flushInternal(IndexDescriptor indexDescriptor) throws IOException
        {
            SegmentWriter writer = index.termType().isLiteral() ? new LiteralIndexWriter(indexDescriptor, index.identifier())
                                                                : new NumericIndexWriter(indexDescriptor, index.identifier(), index.termType().fixedSizeOf());

            return writer.writeCompleteSegment(segmentTrieBuffer.iterator());
        }

        @Override
        public boolean isEmpty()
        {
            return segmentTrieBuffer.numRows() == 0;
        }
    }

    public static class VectorSegmentBuilder extends SegmentBuilder
    {
        private final OnHeapGraph<Integer> graphIndex;

        public VectorSegmentBuilder(StorageAttachedIndex index, NamedMemoryLimiter limiter)
        {
            super(index, limiter);
            graphIndex = new OnHeapGraph<>(index.termType().indexType(), index.indexWriterConfig(), null);
        }

        @Override
        public boolean isEmpty()
        {
            return graphIndex.isEmpty();
        }

        @Override
        protected long addInternal(ByteBuffer term, int segmentRowId)
        {
            return graphIndex.add(term, segmentRowId, OnHeapGraph.InvalidVectorBehavior.IGNORE);
        }

        @Override
        protected SegmentMetadata.ComponentMetadataMap flushInternal(IndexDescriptor indexDescriptor) throws IOException
        {
            return graphIndex.writeData(indexDescriptor, index.identifier(), p -> p);
        }
    }

    public static int getActiveBuilderCount()
    {
        return ACTIVE_BUILDER_COUNT.get();
    }

    private SegmentBuilder(StorageAttachedIndex index, NamedMemoryLimiter limiter)
    {
        this.index = index;
        this.limiter = limiter;
        lastValidSegmentRowID = testLastValidSegmentRowId >= 0 ? testLastValidSegmentRowId : LAST_VALID_SEGMENT_ROW_ID;

        minimumFlushBytes = limiter.limitBytes() / ACTIVE_BUILDER_COUNT.incrementAndGet();
    }

    public SegmentMetadata flush(IndexDescriptor indexDescriptor) throws IOException
    {
        assert !flushed : "Cannot flush an already flushed segment";
        flushed = true;

        if (getRowCount() == 0)
        {
            logger.warn(index.identifier().logMessage("No rows to index during flush of SSTable {}."), indexDescriptor.sstableDescriptor);
            return null;
        }

        SegmentMetadata.ComponentMetadataMap indexMetas = flushInternal(indexDescriptor);

        return new SegmentMetadata(segmentRowIdOffset, rowCount, minSSTableRowId, maxSSTableRowId, minKey, maxKey, minTerm, maxTerm, indexMetas);
    }

    public long add(ByteBuffer term, PrimaryKey key, long sstableRowId)
    {
        assert !flushed : "Cannot add to a flushed segment.";
        assert sstableRowId >= maxSSTableRowId;
        minSSTableRowId = minSSTableRowId < 0 ? sstableRowId : minSSTableRowId;
        maxSSTableRowId = sstableRowId;

        assert maxKey == null || maxKey.compareTo(key) <= 0;
        if (minKey == null)
            minKey = key;
        maxKey = key;

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

        if (rowCount == 0)
        {
            // use first global rowId in the segment as segment rowId offset
            segmentRowIdOffset = sstableRowId;
        }

        rowCount++;

        // segmentRowIdOffset should encode sstableRowId into Integer
        int segmentRowId = castToSegmentRowId(sstableRowId, segmentRowIdOffset);
        maxSegmentRowId = Math.max(maxSegmentRowId, segmentRowId);

        long bytesAllocated = addInternal(term, segmentRowId);
        totalBytesAllocated += bytesAllocated;

        return bytesAllocated;
    }

    public static int castToSegmentRowId(long sstableRowId, long segmentRowIdOffset)
    {
        return Math.toIntExact(sstableRowId - segmentRowIdOffset);
    }

    public long totalBytesAllocated()
    {
        return totalBytesAllocated;
    }

    public boolean hasReachedMinimumFlushSize()
    {
        return totalBytesAllocated >= minimumFlushBytes;
    }

    public long getMinimumFlushBytes()
    {
        return minimumFlushBytes;
    }

    /**
     * This method does three things:
     * <p>
     * 1. It decrements active builder count and updates the global minimum flush size to reflect that.
     * 2. It releases the builder's memory against its limiter.
     * 3. It defensively marks the builder inactive to make sure nothing bad happens if we try to close it twice.
     *
     * @return the number of bytes used by the memory limiter after releasing this builder
     */
    public long release()
    {
        if (active)
        {
            minimumFlushBytes = limiter.limitBytes() / ACTIVE_BUILDER_COUNT.getAndDecrement();
            long used = limiter.decrement(totalBytesAllocated);
            active = false;
            return used;
        }

        logger.warn(index.identifier().logMessage("Attempted to release storage-attached index segment builder memory after builder marked inactive."));
        return limiter.currentBytesUsed();
    }

    public abstract boolean isEmpty();

    protected abstract long addInternal(ByteBuffer term, int segmentRowId);

    protected abstract SegmentMetadata.ComponentMetadataMap flushInternal(IndexDescriptor indexDescriptor) throws IOException;

    public int getRowCount()
    {
        return rowCount;
    }

    /**
     * @return true if next SSTable row ID exceeds max segment row ID
     */
    public boolean exceedsSegmentLimit(long ssTableRowId)
    {
        if (getRowCount() == 0)
            return false;

        // To handle the case where there are many non-indexable rows. eg. rowId-1 and rowId-3B are indexable,
        // the rest are non-indexable. We should flush them as 2 separate segments, because rowId-3B is going
        // to cause error in on-disk index structure with 2B limitation.
        return ssTableRowId - segmentRowIdOffset > lastValidSegmentRowID;
    }

    @VisibleForTesting
    public static void updateLastValidSegmentRowId(long lastValidSegmentRowID)
    {
        testLastValidSegmentRowId = lastValidSegmentRowID;
    }
}
