/*
 * 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.db.virtual;

import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Consumer;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;

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

import org.apache.cassandra.cql3.Operator;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.DataRange;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.DeletionTime;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.Slice;
import org.apache.cassandra.db.Slices;
import org.apache.cassandra.db.context.CounterContext;
import org.apache.cassandra.db.filter.ClusteringIndexFilter;
import org.apache.cassandra.db.filter.ColumnFilter;
import org.apache.cassandra.db.filter.DataLimits;
import org.apache.cassandra.db.filter.RowFilter;
import org.apache.cassandra.db.marshal.CompositeType;
import org.apache.cassandra.db.marshal.CounterColumnType;
import org.apache.cassandra.db.marshal.IntegerType;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.db.partitions.PartitionUpdate;
import org.apache.cassandra.db.partitions.SingletonUnfilteredPartitionIterator;
import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
import org.apache.cassandra.db.rows.AbstractUnfilteredRowIterator;
import org.apache.cassandra.db.rows.BTreeRow;
import org.apache.cassandra.db.rows.BufferCell;
import org.apache.cassandra.db.rows.Cell;
import org.apache.cassandra.db.rows.EncodingStats;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.db.rows.Rows;
import org.apache.cassandra.db.rows.Unfiltered;
import org.apache.cassandra.db.rows.UnfilteredRowIterator;
import org.apache.cassandra.db.rows.UnfilteredRowIterators;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.dht.Bounds;
import org.apache.cassandra.dht.LocalPartitioner;
import org.apache.cassandra.dht.Range;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.exceptions.InvalidRequestException;
import org.apache.cassandra.io.sstable.KeyReader;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.schema.KeyspaceMetadata;
import org.apache.cassandra.schema.Schema;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.serializers.MarshalException;

import static org.apache.cassandra.cql3.statements.RequestValidations.invalidRequest;

/**
 * A virtual table for querying partition keys of SSTables in a specific keyspace.
 *
 * <p>This table is implemented as a virtual table in Cassandra, meaning it does not
 * store data persistently on disk but instead derives its data from live metadata.
 *
 * <p>The CQL equivalent of this virtual table is:
 * <pre>
 * CREATE TABLE system_views.partition_key_statistics (
 *     keyspace_name TEXT,
 *     table_name TEXT,
 *     token_value INT,
 *     key TEXT,
 *     size_estimate COUNTER,
 *     sstables COUNTER,
 *     PRIMARY KEY ((keyspace_name, table_name), token_value, key)
 * );
 * </pre>
 *
 * <p>Note:
 * <ul>
 *     <li>The `size_estimate` and `sstables` columns represent aggregate information about SSTable sizes and counts.</li>
 *     <li>Range queries across multiple tables and updates are not supported as this is a read-only table.</li>
 * </ul>
 */
public class PartitionKeyStatsTable implements VirtualTable
{
    private static final Logger logger = LoggerFactory.getLogger(PartitionKeyStatsTable.class);
    public static final String NAME = "partition_key_statistics";

    private static final String TABLE_READ_ONLY_ERROR = "The specified table is read-only.";
    private static final String UNSUPPORTED_RANGE_QUERY_ERROR = "Range queries are not supported. Please provide both a keyspace and a table name.";
    private static final String REVERSED_QUERY_ERROR = "Reversed queries are not supported.";
    private static final String KEYSPACE_NOT_EXIST_ERROR = "The keyspace '%s' does not exist.";
    private static final String TABLE_NOT_EXIST_ERROR = "The table '%s' does not exist in the keyspace '%s'.";
    private static final String KEY_ONLY_EQUALS_ERROR = "The 'key' column can only be used in an equality query for this virtual table.";
    private static final String KEY_NOT_WITHIN_BOUNDS_ERROR = "The specified 'key' is not within the provided token value bounds.";
    private static final String PARTITIONER_NOT_SUPPORTED = "Partitioner '%s' for table '%s' in keyspace '%s' is not supported.";

    private static final String COLUMN_KEYSPACE_NAME = "keyspace_name";
    private static final String COLUMN_TABLE_NAME = "table_name";
    private static final String COLUMN_TOKEN_VALUE = "token_value";
    private static final String COLUMN_KEY = "key";
    private static final String COLUMN_SIZE_ESTIMATE = "size_estimate";
    private static final String COLUMN_SSTABLES = "sstables";

    private final TableMetadata metadata;
    private final ColumnMetadata sizeEstimateColumn;
    private final ColumnMetadata sstablesColumn;

    @VisibleForTesting
    final CopyOnWriteArrayList<Consumer<DecoratedKey>> readListener = new CopyOnWriteArrayList<>();

    public PartitionKeyStatsTable(String keyspace)
    {
        this.metadata = TableMetadata.builder(keyspace, NAME)
                                     .kind(TableMetadata.Kind.VIRTUAL)
                                     .partitioner(new LocalPartitioner(CompositeType.getInstance(UTF8Type.instance, UTF8Type.instance)))
                                     .addPartitionKeyColumn(COLUMN_KEYSPACE_NAME, UTF8Type.instance)
                                     .addPartitionKeyColumn(COLUMN_TABLE_NAME, UTF8Type.instance)
                                     .addClusteringColumn(COLUMN_TOKEN_VALUE, IntegerType.instance)
                                     .addClusteringColumn(COLUMN_KEY, UTF8Type.instance)
                                     .addRegularColumn(COLUMN_SIZE_ESTIMATE, CounterColumnType.instance)
                                     .addRegularColumn(COLUMN_SSTABLES, CounterColumnType.instance)
                                     .build();
        sizeEstimateColumn = metadata.regularColumns().getSimple(0);
        sstablesColumn = metadata.regularColumns().getSimple(1);
    }

    @Override
    public UnfilteredPartitionIterator select(DecoratedKey partitionKey, ClusteringIndexFilter clusteringIndexFilter, ColumnFilter columnFilter, RowFilter rowFilter, DataLimits limits)
    {
        if (clusteringIndexFilter.isReversed())
            throw new InvalidRequestException(REVERSED_QUERY_ERROR);

        ByteBuffer[] key = ((CompositeType) this.metadata.partitionKeyType).split(partitionKey.getKey());
        String keyspace = UTF8Type.instance.getString(key[0]);
        String table = UTF8Type.instance.getString(key[1]);

        KeyspaceMetadata ksm = Schema.instance.getKeyspaceMetadata(keyspace);
        if (ksm == null)
            throw invalidRequest(KEYSPACE_NOT_EXIST_ERROR, keyspace);

        TableMetadata metadata = ksm.getTableOrViewNullable(table);
        if (metadata == null)
            throw invalidRequest(TABLE_NOT_EXIST_ERROR, table, keyspace);

        if (!metadata.partitioner.supportsSplitting())
            throw invalidRequest(PARTITIONER_NOT_SUPPORTED, metadata.partitioner.getClass().getName(), table, keyspace);

        AbstractBounds<PartitionPosition> range = getBounds(metadata, clusteringIndexFilter, rowFilter);
        return new SingletonUnfilteredPartitionIterator(select(partitionKey, metadata, clusteringIndexFilter, range));
    }

    private List<SSTableReader> getSStables(TableMetadata metadata, AbstractBounds<PartitionPosition> range)
    {
        return Lists.newArrayList(ColumnFamilyStore.getIfExists(metadata).getTracker().getView().liveSSTablesInBounds(range.left, range.right));
    }

    private UnfilteredRowIterator select(DecoratedKey partitionKey, TableMetadata metadata, ClusteringIndexFilter clusteringIndexFilter, AbstractBounds<PartitionPosition> range)
    {
        List<SSTableReader> sstables = getSStables(metadata, range);
        if (sstables.isEmpty())
            return UnfilteredRowIterators.noRowsIterator(metadata, partitionKey, Rows.EMPTY_STATIC_ROW, DeletionTime.LIVE, false);

        List<UnfilteredRowIterator> sstableIterators = Lists.newArrayList();
        for (SSTableReader sstable : sstables)
            sstableIterators.add(getSStableRowIterator(metadata, partitionKey, sstable, clusteringIndexFilter, range));

        return UnfilteredRowIterators.merge(sstableIterators);
    }

    private UnfilteredRowIterator getSStableRowIterator(TableMetadata target, DecoratedKey partitionKey, SSTableReader sstable, ClusteringIndexFilter filter, AbstractBounds<PartitionPosition> range)
    {
        final KeyReader reader;
        try
        {
            // ignore warning on try-with-resources, the reader will be closed on endOfData or close
            reader = sstable.keyReader(range.left);
        }
        catch (IOException e)
        {
            logger.error("Error generating keyReader for SSTable: {}", sstable, e);
            throw new RuntimeException(e);
        }

        return new AbstractUnfilteredRowIterator(metadata, partitionKey, DeletionTime.LIVE,
                                                 metadata.regularAndStaticColumns(), Rows.EMPTY_STATIC_ROW,
                                                 false, EncodingStats.NO_STATS)
        {
            public Unfiltered endOfData()
            {
                reader.close();
                return super.endOfData();
            }

            public void close()
            {
                reader.close();
            }

            private Row buildRow(Clustering<?> clustering, long size)
            {
                Row.Builder row = BTreeRow.sortedBuilder();
                row.newRow(clustering);
                row.addCell(cell(sizeEstimateColumn, CounterContext.instance().createUpdate(size)));
                row.addCell(cell(sstablesColumn, CounterContext.instance().createUpdate(1)));
                return row.build();
            }

            @Override
            protected Unfiltered computeNext()
            {
                while (!reader.isExhausted())
                {
                    DecoratedKey key = target.partitioner.decorateKey(reader.key());

                    for (Consumer<DecoratedKey> listener : readListener)
                        listener.accept(key);

                    // Store the reader's current data position to calculate size later
                    long lastPosition = reader.dataPosition();
                    try
                    {
                        // Advance the reader to the next key for the next iteration. Also by moving to next key
                        // we move the dataPosition to the start of the next key for calculating size
                        reader.advance();
                    }
                    catch (IOException e)
                    {
                        logger.error("Error advancing reader for SSTable: {}", sstable, e);
                        return endOfData();
                    }

                    // Calculate the size of the current key. If EOF use the length of the file
                    long current = reader.dataPosition() == -1 ? sstable.uncompressedLength() : reader.dataPosition();
                    long size = current - lastPosition;

                    String keyString = target.partitionKeyType.getString(key.getKey());

                    // Check if the current key is outside the queried range; if so, stop
                    if (range.right.compareTo(key) < 0)
                        return endOfData();

                    // Convert the token to a string and create a clustering object
                    String tokenString = key.getToken().toString();
                    Clustering<?> clustering = Clustering.make(
                    IntegerType.instance.decompose(new BigInteger(tokenString)),
                    UTF8Type.instance.decompose(keyString)
                    );

                    // Check if the current clustering matches the filter; if so, return the row
                    if (filter.selects(clustering))
                        return buildRow(clustering, size);
                }
                return endOfData();
            }
        };
    }

    /**
     * This converts the clustering token/key into the partition level token/key for the target table. Also provides an
     * optimization from RowFilter when a `key` is specified with or without the clustering `token` being set.
     */
    private AbstractBounds<PartitionPosition> getBounds(TableMetadata target, ClusteringIndexFilter clusteringIndexFilter, RowFilter rowFilter)
    {
        Slices s = clusteringIndexFilter.getSlices(target);
        Token startToken = target.partitioner.getMinimumToken();
        Token endToken = target.partitioner.getMaximumTokenForSplitting();
        BigInteger startTokenValue = new BigInteger(endToken.getTokenValue().toString(), 10);
        BigInteger endTokenValue = new BigInteger(startToken.getTokenValue().toString(), 10);

        // find min/max token values from the clustering key
        for (int i = 0; i < s.size(); i++)
        {
            Slice slice = s.get(i);
            if (!slice.start().isEmpty())
            {
                startTokenValue = startTokenValue.min(IntegerType.instance.compose(slice.start().bufferAt(0)));
                startToken = target.partitioner.getTokenFactory().fromString(startTokenValue.toString());
            }
            if (!slice.end().isEmpty())
            {
                endTokenValue = endTokenValue.max(IntegerType.instance.compose(slice.end().bufferAt(0)));
                endToken = target.partitioner.getTokenFactory().fromString(endTokenValue.toString());
            }
        }

        // override min/max of token if the `key` is specified
        for (RowFilter.Expression expression : rowFilter.getExpressions())
        {
            if (expression.column().name.toString().equals(COLUMN_KEY))
            {
                if (expression.operator() != Operator.EQ)
                    throw new InvalidRequestException(KEY_ONLY_EQUALS_ERROR);

                String keyString = UTF8Type.instance.compose(expression.getIndexValue());
                ByteBuffer keyAsBB;
                try
                {
                    keyAsBB = target.partitionKeyType.fromString(keyString);
                }
                catch (MarshalException ex)
                {
                    throw new InvalidRequestException(ex.getMessage());
                }
                DecoratedKey decoratedKey = target.partitioner.decorateKey(keyAsBB);

                if (!DataRange.forKeyRange(new Range<>(startToken.minKeyBound(), endToken.maxKeyBound())).contains(decoratedKey.getToken().minKeyBound()))
                    throw new InvalidRequestException(KEY_NOT_WITHIN_BOUNDS_ERROR);

                return Bounds.bounds(decoratedKey, true, decoratedKey, true);
            }
        }
        return Bounds.bounds(startToken.minKeyBound(), true, endToken.maxKeyBound(), true);
    }

    private static Cell<?> cell(ColumnMetadata column, ByteBuffer value)
    {
        return BufferCell.live(column, 1L, value);
    }

    @Override
    public TableMetadata metadata()
    {
        return this.metadata;
    }

    @Override
    public UnfilteredPartitionIterator select(DataRange dataRange, ColumnFilter columnFilter, RowFilter rowFilter, DataLimits limits)
    {
        throw new InvalidRequestException(UNSUPPORTED_RANGE_QUERY_ERROR);
    }

    @Override
    public void truncate()
    {
        throw new InvalidRequestException(TABLE_READ_ONLY_ERROR);
    }

    @Override
    public void apply(PartitionUpdate update)
    {
        throw new InvalidRequestException(TABLE_READ_ONLY_ERROR);
    }
}
