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

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.PartitionPosition;
import org.apache.cassandra.db.memtable.Memtable;
import org.apache.cassandra.dht.AbstractBounds;
import org.apache.cassandra.index.sai.disk.SSTableIndex;
import org.apache.cassandra.index.sai.memory.MemtableIndex;
import org.apache.cassandra.index.sai.view.View;
import org.apache.cassandra.io.sstable.format.SSTableReader;

/**
 * Build a query specific view of the on-disk indexes for a query. This will return a
 * {@link Collection} of {@link Expression} and {@link SSTableIndex}s that represent
 * the on-disk data for a query.
 * <p>
 * The query view will include all the indexed expressions even if they don't have any
 * on-disk data. This in necessary because the query view is used to query in-memory
 * data as well as the attached on-disk indexes.
 */
public class QueryViewBuilder
{
    private final Collection<Expression> expressions;
    private final AbstractBounds<PartitionPosition> range;

    QueryViewBuilder(Collection<Expression> expressions, AbstractBounds<PartitionPosition> range)
    {
        this.expressions = expressions;
        this.range = range;
    }

    public static class QueryExpressionView
    {
        public final Expression expression;
        public final Collection<MemtableIndex> memtableIndexes;
        public final Collection<SSTableIndex> sstableIndexes;

        public QueryExpressionView(Expression expression, Collection<MemtableIndex> memtableIndexes, Collection<SSTableIndex> sstableIndexes)
        {
            this.expression = expression;
            this.memtableIndexes = memtableIndexes;
            this.sstableIndexes = sstableIndexes;
        }

        public ColumnFamilyStore.ViewFragment computeViewFragment()
        {
            // Because the SSTableIndex holds a reference to the SSTableReader, we know the sstable is still accessible
            // so it is safe to build a view fragment.
            List<Memtable> memtables = memtableIndexes.stream().map(MemtableIndex::getMemtable).collect(Collectors.toList());
            List<SSTableReader> sstableReaders = sstableIndexes.stream().map(SSTableIndex::getSSTable).collect(Collectors.toList());
            return new ColumnFamilyStore.ViewFragment(sstableReaders, memtables);
        }
    }

    public static class QueryView implements AutoCloseable
    {
        public final Collection<QueryExpressionView> view;
        public final Set<SSTableIndex> referencedIndexes;

        public QueryView(Collection<QueryExpressionView> view, Set<SSTableIndex> referencedIndexes)
        {
            this.view = view;
            this.referencedIndexes = referencedIndexes;
        }

        @Override
        public void close()
        {
            referencedIndexes.forEach(SSTableIndex::releaseQuietly);
            referencedIndexes.clear();
        }
    }

    protected QueryView build()
    {
        Set<SSTableIndex> referencedIndexes = new HashSet<>();
        while (true)
        {
            referencedIndexes.clear();
            boolean failed = false;

            Collection<QueryExpressionView> view = getQueryView(expressions);
            for (SSTableIndex index : view.stream().map(v -> v.sstableIndexes).flatMap(Collection::stream).collect(Collectors.toList()))
            {
                if (index.reference())
                    referencedIndexes.add(index);
                else
                    failed = true;
            }

            if (failed)
                referencedIndexes.forEach(SSTableIndex::release);
            else
                return new QueryView(view, referencedIndexes);
        }
    }

    private Collection<QueryExpressionView> getQueryView(Collection<Expression> expressions)
    {
        List<QueryExpressionView> queryView = new ArrayList<>();

        for (Expression expression : expressions)
        {
            // Non-index column query should only act as FILTER BY for satisfiedBy(Row) method
            // because otherwise it likely to go through the whole index.
            if (expression.isNotIndexed())
                continue;

            // Fetch the memtables first to ensure we don't miss any newly flushed memtable index
            Collection<MemtableIndex> memtableIndexes = expression.getIndex().memtableIndexManager().getLiveMemtableIndexesSnapshot();
            // Select all the sstable indexes that have a term range that is satisfied by this expression and
            // overlap with the key range being queried.
            View view = expression.getIndex().view();
            Collection<SSTableIndex> sstableIndexes = selectIndexesInRange(view.match(expression));
            queryView.add(new QueryExpressionView(expression, memtableIndexes, sstableIndexes));
        }

        return queryView;
    }

    private List<SSTableIndex> selectIndexesInRange(Collection<SSTableIndex> indexes)
    {
        return indexes.stream().filter(this::indexInRange).sorted(SSTableIndex.COMPARATOR).collect(Collectors.toList());
    }

    private boolean indexInRange(SSTableIndex index)
    {
        SSTableReader sstable = index.getSSTable();
        return range.left.compareTo(sstable.getLast()) <= 0 && (range.right.isMinimum() || sstable.getFirst().compareTo(range.right) <= 0);
    }
}
