/*
 * 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;

import java.io.IOException;
import java.util.Collections;
import java.util.Iterator;
import java.util.concurrent.TimeUnit;

import com.google.common.base.Stopwatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.carrotsearch.hppc.LongArrayList;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.index.sai.disk.PerColumnIndexWriter;
import org.apache.cassandra.index.sai.disk.RowMapping;
import org.apache.cassandra.index.sai.disk.format.IndexComponent;
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.segment.SegmentMetadata;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentWriter;
import org.apache.cassandra.index.sai.disk.v1.trie.LiteralIndexWriter;
import org.apache.cassandra.index.sai.memory.MemtableIndex;
import org.apache.cassandra.index.sai.memory.MemtableTermsIterator;
import org.apache.cassandra.index.sai.metrics.IndexMetrics;
import org.apache.cassandra.index.sai.utils.IndexIdentifier;
import org.apache.cassandra.index.sai.utils.IndexTermType;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;

/**
 * Column index writer that flushes indexed data directly from the corresponding Memtable index, without buffering index
 * data in memory.
 */
public class MemtableIndexWriter implements PerColumnIndexWriter
{
    private static final Logger logger = LoggerFactory.getLogger(MemtableIndexWriter.class);
    private static final int NO_ROWS = -1;

    private final IndexDescriptor indexDescriptor;
    private final IndexTermType indexTermType;
    private final IndexIdentifier indexIdentifier;
    private final IndexMetrics indexMetrics;
    private final MemtableIndex memtable;
    private final RowMapping rowMapping;

    private PrimaryKey minKey;
    private PrimaryKey maxKey;
    private long maxSSTableRowId = NO_ROWS;
    private int rowCount;

    public MemtableIndexWriter(MemtableIndex memtable,
                               IndexDescriptor indexDescriptor,
                               IndexTermType indexTermType,
                               IndexIdentifier indexIdentifier,
                               IndexMetrics indexMetrics,
                               RowMapping rowMapping)
    {
        assert rowMapping != null && rowMapping != RowMapping.DUMMY : "Row mapping must exist during FLUSH.";

        this.indexDescriptor = indexDescriptor;
        this.indexTermType = indexTermType;
        this.indexIdentifier = indexIdentifier;
        this.indexMetrics = indexMetrics;
        this.memtable = memtable;
        this.rowMapping = rowMapping;
    }

    @Override
    public void addRow(PrimaryKey key, Row row, long sstableRowId)
    {
        // Memtable indexes are flushed directly to disk with the aid of a mapping between primary
        // keys and row IDs in the flushing SSTable. This writer, therefore, does nothing in
        // response to the flushing of individual rows except for keeping index-specific statistics.
        boolean isStatic = indexTermType.columnMetadata().isStatic();
        boolean isPartitionKey = indexTermType.columnMetadata().isPartitionKey();

        // Indexes on static columns should only track static rows, and indexes on non-static columns 
        // should only track non-static rows. (Within a partition, the row ID for a static row will always
        // come before any non-static row.) The only exception to this is indexes on partition key elements.
        if ((key.kind() == PrimaryKey.Kind.STATIC && (isStatic || isPartitionKey)) || key.kind() != PrimaryKey.Kind.STATIC && !isStatic)
        {
            if (minKey == null)
                minKey = key;

            maxKey = key;
            rowCount++;
            maxSSTableRowId = Math.max(maxSSTableRowId, sstableRowId);
        }
    }

    @Override
    public void abort(Throwable cause)
    {
        if (cause == null)
            // This commonly occurs when a Memtable has no rows to flush, and is harmless:
            logger.debug(indexIdentifier.logMessage("Aborting index memtable flush for {}..."), indexDescriptor.sstableDescriptor);
        else
            logger.warn(indexIdentifier.logMessage("Aborting index memtable flush for {}..."), indexDescriptor.sstableDescriptor, cause);

        indexDescriptor.deleteColumnIndex(indexTermType, indexIdentifier);
    }

    @Override
    public void complete(Stopwatch stopwatch) throws IOException
    {
        assert rowMapping.isComplete() : "Cannot complete the memtable index writer because the row mapping is not complete";

        long start = stopwatch.elapsed(TimeUnit.MILLISECONDS);

        try
        {
            if (maxSSTableRowId == -1 || memtable == null || memtable.isEmpty())
            {
                logger.debug(indexIdentifier.logMessage("No indexed rows to flush from SSTable {}."), indexDescriptor.sstableDescriptor);
                // Write a completion marker even though we haven't written anything to the index,
                // so we won't try to build the index again for the SSTable
                ColumnCompletionMarkerUtil.create(indexDescriptor, indexIdentifier, true);

                return;
            }

            if (indexTermType.isVector())
            {
                flushVectorIndex(start, stopwatch);
            }
            else
            {
                final Iterator<Pair<ByteComparable, LongArrayList>> iterator = rowMapping.merge(memtable);

                long cellCount = 0;
                if (iterator.hasNext())
                {
                    try (MemtableTermsIterator terms = new MemtableTermsIterator(memtable.getMinTerm(), memtable.getMaxTerm(), iterator))
                    {
                        cellCount = flush(terms);
                    }
                }
                completeIndexFlush(cellCount, start, stopwatch);
            }
        }
        catch (Throwable t)
        {
            logger.error(indexIdentifier.logMessage("Error while flushing index {}"), t.getMessage(), t);
            indexMetrics.memtableIndexFlushErrors.inc();

            throw t;
        }
    }

    @Override
    public void onSSTableWriterSwitched(Stopwatch stopwatch) throws IOException
    {
        // no-op for memtable index where all terms are already inside memory index, we can't get rid of memory index
        // until full flush are completed
    }

    private long flush(MemtableTermsIterator terms) throws IOException
    {
        SegmentWriter writer = indexTermType.isLiteral() ? new LiteralIndexWriter(indexDescriptor, indexIdentifier)
                                                         : new NumericIndexWriter(indexDescriptor,
                                                                                  indexIdentifier,
                                                                                  indexTermType.fixedSizeOf());

        SegmentMetadata.ComponentMetadataMap indexMetas = writer.writeCompleteSegment(terms);
        long numRows = writer.getNumberOfRows();

        // If no rows were written we need to delete any created column index components
        // so that the index is correctly identified as being empty (only having a completion marker)
        if (numRows == 0)
        {
            indexDescriptor.deleteColumnIndex(indexTermType, indexIdentifier);
            return 0;
        }

        // During index memtable flush, the data is sorted based on terms.
        SegmentMetadata metadata = new SegmentMetadata(0,
                                                       numRows,
                                                       terms.getMinSSTableRowId(), terms.getMaxSSTableRowId(),
                                                       minKey, maxKey, 
                                                       terms.getMinTerm(), terms.getMaxTerm(),
                                                       indexMetas);

        try (MetadataWriter metadataWriter = new MetadataWriter(indexDescriptor.openPerIndexOutput(IndexComponent.META, indexIdentifier)))
        {
            SegmentMetadata.write(metadataWriter, Collections.singletonList(metadata));
        }

        return numRows;
    }

    private void flushVectorIndex(long startTime, Stopwatch stopwatch) throws IOException
    {
        SegmentMetadata.ComponentMetadataMap metadataMap = memtable.writeDirect(indexDescriptor, indexIdentifier, rowMapping::get);
        completeIndexFlush(rowCount, startTime, stopwatch);

        SegmentMetadata metadata = new SegmentMetadata(0,
                                                       rowCount,
                                                       0, maxSSTableRowId,
                                                       minKey, maxKey, 
                                                       ByteBufferUtil.bytes(0), ByteBufferUtil.bytes(0),
                                                       metadataMap);

        try (MetadataWriter writer = new MetadataWriter(indexDescriptor.openPerIndexOutput(IndexComponent.META, indexIdentifier)))
        {
            SegmentMetadata.write(writer, Collections.singletonList(metadata));
        }
    }

    private void completeIndexFlush(long cellCount, long startTime, Stopwatch stopwatch) throws IOException
    {
        // create a completion marker indicating that the index is complete
        ColumnCompletionMarkerUtil.create(indexDescriptor, indexIdentifier, cellCount == 0);

        indexMetrics.memtableIndexFlushCount.inc();

        long elapsedTime = stopwatch.elapsed(TimeUnit.MILLISECONDS);

        logger.debug(indexIdentifier.logMessage("Completed flushing {} memtable index cells to SSTable {}. Duration: {} ms. Total elapsed: {} ms"),
                     cellCount,
                     indexDescriptor.sstableDescriptor,
                     elapsedTime - startTime,
                     elapsedTime);

        indexMetrics.memtableFlushCellsPerSecond.update((long) (cellCount * 1000.0 / Math.max(1, elapsedTime - startTime)));
    }
}
