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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.function.Supplier;
import java.util.function.UnaryOperator;

import javax.annotation.Nullable;

import com.google.common.base.Joiner;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.DeletionTime;
import org.apache.cassandra.db.ReadCommand;
import org.apache.cassandra.db.ReadResponse;
import org.apache.cassandra.db.filter.DataLimits;
import org.apache.cassandra.db.partitions.PartitionIterator;
import org.apache.cassandra.db.partitions.PartitionIterators;
import org.apache.cassandra.db.partitions.UnfilteredPartitionIterator;
import org.apache.cassandra.db.partitions.UnfilteredPartitionIterators;
import org.apache.cassandra.db.rows.RangeTombstoneMarker;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.db.rows.UnfilteredRowIterator;
import org.apache.cassandra.db.rows.UnfilteredRowIterators;
import org.apache.cassandra.db.transform.EmptyPartitionsDiscarder;
import org.apache.cassandra.db.transform.Filter;
import org.apache.cassandra.db.transform.FilteredPartitions;
import org.apache.cassandra.db.transform.Transformation;
import org.apache.cassandra.index.Index;
import org.apache.cassandra.locator.Endpoints;
import org.apache.cassandra.locator.ReplicaPlan;
import org.apache.cassandra.net.Message;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.service.reads.repair.NoopReadRepair;
import org.apache.cassandra.service.reads.repair.ReadRepair;
import org.apache.cassandra.service.reads.repair.RepairedDataTracker;
import org.apache.cassandra.service.reads.repair.RepairedDataVerifier;
import org.apache.cassandra.transport.Dispatcher;

import static com.google.common.collect.Iterables.*;

public class DataResolver<E extends Endpoints<E>, P extends ReplicaPlan.ForRead<E, P>> extends ResponseResolver<E, P>
{
    private final boolean enforceStrictLiveness;
    private final ReadRepair<E, P> readRepair;
    private final boolean trackRepairedStatus;

    public DataResolver(ReadCommand command, Supplier<? extends P> replicaPlan, ReadRepair<E, P> readRepair, Dispatcher.RequestTime requestTime)
    {
        this(command, replicaPlan, readRepair, requestTime, false);
    }

    public DataResolver(ReadCommand command, Supplier<? extends P> replicaPlan, ReadRepair<E, P> readRepair, Dispatcher.RequestTime requestTime, boolean trackRepairedStatus)
    {
        super(command, replicaPlan, requestTime);
        this.enforceStrictLiveness = command.metadata().enforceStrictLiveness();
        this.readRepair = readRepair;
        this.trackRepairedStatus = trackRepairedStatus;
    }

    public PartitionIterator getData()
    {
        ReadResponse response = responses.get(0).payload;
        return UnfilteredPartitionIterators.filter(response.makeIterator(command), command.nowInSec());
    }

    public boolean isDataPresent()
    {
        return !responses.isEmpty();
    }

    public PartitionIterator resolve()
    {
        return resolve(null);
    }

    public PartitionIterator resolve(@Nullable Runnable runOnShortRead)
    {
        // We could get more responses while this method runs, which is ok (we're happy to ignore any response not here
        // at the beginning of this method), so grab the response count once and use that through the method.
        Collection<Message<ReadResponse>> messages = responses.snapshot();
        assert !any(messages, msg -> msg.payload.isDigestResponse());

        E replicas = replicaPlan().readCandidates().select(transform(messages, Message::from), false);

        // If requested, inspect each response for a digest of the replica's repaired data set
        RepairedDataTracker repairedDataTracker = trackRepairedStatus
                                                  ? new RepairedDataTracker(getRepairedDataVerifier(command))
                                                  : null;
        if (repairedDataTracker != null)
        {
            messages.forEach(msg -> {
                if (msg.payload.mayIncludeRepairedDigest() && replicas.byEndpoint().get(msg.from()).isFull())
                {
                    repairedDataTracker.recordDigest(msg.from(),
                                                     msg.payload.repairedDataDigest(),
                                                     msg.payload.isRepairedDigestConclusive());
                }
            });
        }

        if (usesReplicaFilteringProtection())
            return resolveWithReplicaFilteringProtection(replicas, repairedDataTracker);

        ResolveContext context = new ResolveContext(replicas, true);
        return resolveWithReadRepair(context,
                                     i -> shortReadProtectedResponse(i, context, runOnShortRead),
                                     UnaryOperator.identity(),
                                     repairedDataTracker);
    }

    private boolean usesReplicaFilteringProtection()
    {
        if (command.rowFilter().isEmpty())
            return false;

        if (command.isTopK())
            return false;

        Index.QueryPlan queryPlan = command.indexQueryPlan();
        if (queryPlan == null)
            return true;

        return queryPlan.supportsReplicaFilteringProtection(command.rowFilter());
    }

    private class ResolveContext
    {
        private final E replicas;
        private final DataLimits.Counter mergedResultCounter;

        /**
         * @param replicas the collection of {@link Endpoints} involved in the query
         * @param enforceLimits whether or not to enforce counter limits in this context
         */
        private ResolveContext(E replicas, boolean enforceLimits)
        {
            this.replicas = replicas;
            this.mergedResultCounter = command.limits().newCounter(command.nowInSec(),
                                                                   true,
                                                                   command.selectsFullPartition(),
                                                                   enforceStrictLiveness);

            // In case of top-k query, do not trim reconciled rows here because QueryPlan#postProcessor()
            // needs to compare all rows. Also avoid enforcing the limit if explicitly requested.
            if (command.isTopK() || !enforceLimits)
                this.mergedResultCounter.onlyCount();
        }

        private boolean needsReadRepair()
        {
            // Each replica may return different estimated top-K rows, it doesn't mean data is not replicated.
            // Even though top-K queries are limited to CL ONE & LOCAL-ONE, they use the ScanAllRangesCommandIterator
            // that combines the separate replica plans of each data range into a single replica plan. This is an
            // optimisation but can result in the number of replicas being > 1.
            if (command.isTopK())
                return false;

            return replicas.size() > 1;
        }

        private boolean needShortReadProtection()
        {
            // SRP doesn't make sense for top-k which needs to re-query replica with larger limit instead of fetching more partitions
            if (command.isTopK())
                return false;

            // If we have only one result, there is no read repair to do, and we can't get short reads
            // Also, so-called "short reads" stems from nodes returning only a subset of the results they have for a
            // partition due to the limit, but that subset not being enough post-reconciliation. So if we don't have limit,
            // don't bother protecting against short reads.
            return replicas.size() > 1 && !command.limits().isUnlimited();
        }
    }

    @FunctionalInterface
    private interface ResponseProvider
    {
        UnfilteredPartitionIterator getResponse(int i);
    }

    private UnfilteredPartitionIterator shortReadProtectedResponse(int i, ResolveContext context, @Nullable Runnable onShortRead)
    {
        UnfilteredPartitionIterator originalResponse = responses.get(i).payload.makeIterator(command);

        return context.needShortReadProtection()
               ? ShortReadProtection.extend(context.replicas.get(i),
                                            () -> { responses.clearUnsafe(i); if (onShortRead != null) onShortRead.run(); },
                                            originalResponse,
                                            command,
                                            context.mergedResultCounter,
                                            requestTime,
                                            enforceStrictLiveness)
               : originalResponse;
    }

    private PartitionIterator resolveWithReadRepair(ResolveContext context,
                                                    ResponseProvider responseProvider,
                                                    UnaryOperator<PartitionIterator> preCountFilter,
                                                    RepairedDataTracker repairedDataTracker)
    {
        UnfilteredPartitionIterators.MergeListener listener = null;
        if (context.needsReadRepair() && readRepair != NoopReadRepair.instance)
        {
            P sources = replicaPlan.get().withContacts(context.replicas);
            listener = wrapMergeListener(readRepair.getMergeListener(sources), sources, repairedDataTracker);
        }

        return resolveInternal(context, listener, responseProvider, preCountFilter);
    }

    private PartitionIterator resolveWithReplicaFilteringProtection(E replicas, RepairedDataTracker repairedDataTracker)
    {
        // Protecting against inconsistent replica filtering (some replica returning a row that is outdated but that
        // wouldn't be removed by normal reconciliation because up-to-date replica have filtered the up-to-date version
        // of that row) involves 3 main elements:
        //   1) We combine an unlimited, short-read-protected iterator of unfiltered partitions with a merge listener 
        //      that identifies potentially "out-of-date" rows. A row is considered out-of-date if its merged form is 
        //      non-empty, and we receive no response from some replica. It is also potentially out-of-date if there is
        //      disagreement around the value of a single column in a non-empty row, and some replica provides no value
        //      for that column. In either case, it is possible that filtering at the "silent" replica has produced a 
        //      more up-to-date result.
        //   2) This iterator is passed to the standard resolution process with read-repair, but is first wrapped in a
        //      response provider that lazily "completes" potentially out-of-date rows by directly querying them on the
        //      replicas that were previously silent. As this iterator is consumed, it caches valid data for potentially
        //      out-of-date rows, and this cached data is merged with the fetched data as rows are requested. Only rows 
        //      in the partition being evalutated will be cached (then released when the partition is consumed).
        //   3) After a "complete" row is materialized, it must pass the row filter supplied by the original query
        //      before it counts against the limit. If this "pre-count" filter causes a short read, additional rows
        //      will be fetched from the first-phase iterator.

        ReplicaFilteringProtection<E> rfp = new ReplicaFilteringProtection<>(replicaPlan().keyspace(),
                                                                             command,
                                                                             replicaPlan().consistencyLevel(),
                                                                             requestTime,
                                                                             replicas,
                                                                             DatabaseDescriptor.getCachedReplicaRowsWarnThreshold(),
                                                                             DatabaseDescriptor.getCachedReplicaRowsFailThreshold());

        ResolveContext firstPhaseContext = new ResolveContext(replicas, false);
        PartitionIterator firstPhasePartitions = resolveInternal(firstPhaseContext,
                                                                 rfp.mergeController(),
                                                                 i -> shortReadProtectedResponse(i, firstPhaseContext, null),
                                                                 null);

        ResolveContext secondPhaseContext = new ResolveContext(replicas, true);
        PartitionIterator completedPartitions = resolveWithReadRepair(secondPhaseContext,
                                                                      i -> rfp.queryProtectedPartitions(firstPhasePartitions, i),
                                                                      preCountFilterForReplicaFilteringProtection(),
                                                                      repairedDataTracker);

        // Ensure that the RFP instance has a chance to record metrics when the iterator closes.
        return PartitionIterators.doOnClose(completedPartitions, firstPhasePartitions::close);
    }

    private  UnaryOperator<PartitionIterator> preCountFilterForReplicaFilteringProtection()
    {
        // Key columns are immutable and should never need to participate in replica filtering
        if (!command.rowFilter().hasNonKeyExpression())
            return results -> results;

        return results -> {
            Index.Searcher searcher = command.indexSearcher();
            // in case of "ALLOW FILTERING" without index
            if (searcher == null)
                return command.rowFilter().filter(results, command.metadata(), command.nowInSec());
            return searcher.filterReplicaFilteringProtection(results);
        };
    }

    private PartitionIterator resolveInternal(ResolveContext context,
                                              UnfilteredPartitionIterators.MergeListener mergeListener,
                                              ResponseProvider responseProvider,
                                              @Nullable UnaryOperator<PartitionIterator> preCountFilter)
    {
        int count = context.replicas.size();
        List<UnfilteredPartitionIterator> results = new ArrayList<>(count);
        for (int i = 0; i < count; i++)
            results.add(responseProvider.getResponse(i));

        /*
         * Even though every response, individually, will honor the limit, it is possible that we will, after the merge,
         * have more rows than the client requested. To make sure that we still conform to the original limit,
         * we apply a top-level post-reconciliation counter to the merged partition iterator.
         *
         * Short read protection logic (ShortReadRowsProtection.moreContents()) relies on this counter to be applied
         * to the current partition to work. For this reason we have to apply the counter transformation before
         * empty partition discard logic kicks in - for it will eagerly consume the iterator.
         *
         * That's why the order here is: 1) merge; 2) filter rows; 3) count; 4) discard empty partitions
         *
         * See CASSANDRA-13747 for more details.
         */

        UnfilteredPartitionIterator merged = UnfilteredPartitionIterators.merge(results, mergeListener);
        Filter filter = new Filter(command.nowInSec(), command.metadata().enforceStrictLiveness());
        FilteredPartitions filtered = FilteredPartitions.filter(merged, filter);

        PartitionIterator counted = preCountFilter == null
                                    ? filtered
                                    : Transformation.apply(preCountFilter.apply(filtered), context.mergedResultCounter);

        return Transformation.apply(counted, new EmptyPartitionsDiscarder());
    }

    protected RepairedDataVerifier getRepairedDataVerifier(ReadCommand command)
    {
        return RepairedDataVerifier.verifier(command);
    }

    private String makeResponsesDebugString(DecoratedKey partitionKey)
    {
        return Joiner.on(",\n").join(transform(getMessages().snapshot(), m -> m.from() + " => " + m.payload.toDebugString(command, partitionKey)));
    }

    private UnfilteredPartitionIterators.MergeListener wrapMergeListener(UnfilteredPartitionIterators.MergeListener partitionListener,
                                                                         P sources,
                                                                         RepairedDataTracker repairedDataTracker)
    {
        // Avoid wrapping no-op listener as it doesn't throw, unless we're tracking repaired status
        // in which case we need to inject the tracker & verify on close
        if (partitionListener == UnfilteredPartitionIterators.MergeListener.NOOP)
        {
            if (repairedDataTracker == null)
                return partitionListener;

            return new UnfilteredPartitionIterators.MergeListener()
            {

                public UnfilteredRowIterators.MergeListener getRowMergeListener(DecoratedKey partitionKey, List<UnfilteredRowIterator> versions)
                {
                    return UnfilteredRowIterators.MergeListener.NOOP;
                }

                public void close()
                {
                    repairedDataTracker.verify();
                }
            };
        }

        return new UnfilteredPartitionIterators.MergeListener()
        {
            public UnfilteredRowIterators.MergeListener getRowMergeListener(DecoratedKey partitionKey, List<UnfilteredRowIterator> versions)
            {
                UnfilteredRowIterators.MergeListener rowListener = partitionListener.getRowMergeListener(partitionKey, versions);

                return new UnfilteredRowIterators.MergeListener()
                {
                    public void onMergedPartitionLevelDeletion(DeletionTime mergedDeletion, DeletionTime[] versions)
                    {
                        try
                        {
                            rowListener.onMergedPartitionLevelDeletion(mergedDeletion, versions);
                        }
                        catch (AssertionError e)
                        {
                            // The following can be pretty verbose, but it's really only triggered if a bug happen, so we'd
                            // rather get more info to debug than not.
                            TableMetadata table = command.metadata();
                            String details = String.format("Error merging partition level deletion on %s: merged=%s, versions=%s, sources={%s}, debug info:%n %s",
                                                           table,
                                                           mergedDeletion == null ? "null" : mergedDeletion.toString(),
                                                           '[' + Joiner.on(", ").join(transform(Arrays.asList(versions), rt -> rt == null ? "null" : rt.toString())) + ']',
                                                           sources.contacts(),
                                                           makeResponsesDebugString(partitionKey));
                            throw new AssertionError(details, e);
                        }
                    }

                    public void onMergedRows(Row merged, Row[] versions)
                    {
                        try
                        {
                            rowListener.onMergedRows(merged, versions);
                        }
                        catch (AssertionError e)
                        {
                            // The following can be pretty verbose, but it's really only triggered if a bug happen, so we'd
                            // rather get more info to debug than not.
                            TableMetadata table = command.metadata();
                            String details = String.format("Error merging rows on %s: merged=%s, versions=%s, sources={%s}, debug info:%n %s",
                                                           table,
                                                           merged == null ? "null" : merged.toString(table),
                                                           '[' + Joiner.on(", ").join(transform(Arrays.asList(versions), rt -> rt == null ? "null" : rt.toString(table))) + ']',
                                                           sources.contacts(),
                                                           makeResponsesDebugString(partitionKey));
                            throw new AssertionError(details, e);
                        }
                    }

                    public void onMergedRangeTombstoneMarkers(RangeTombstoneMarker merged, RangeTombstoneMarker[] versions)
                    {
                        try
                        {
                            // The code for merging range tombstones is a tad complex, and we had the assertions there triggered
                            // unexpectedly in a few occasions (CASSANDRA-13237, CASSANDRA-13719). It's hard to get insights
                            // when that happen without more context that what the assertion errors give us however, hence the
                            // catch here that basically gather as much as context as reasonable.
                            rowListener.onMergedRangeTombstoneMarkers(merged, versions);
                        }
                        catch (AssertionError e)
                        {

                            // The following can be pretty verbose, but it's really only triggered if a bug happen, so we'd
                            // rather get more info to debug than not.
                            TableMetadata table = command.metadata();
                            String details = String.format("Error merging RTs on %s: merged=%s, versions=%s, sources={%s}, debug info:%n %s",
                                                           table,
                                                           merged == null ? "null" : merged.toString(table),
                                                           '[' + Joiner.on(", ").join(transform(Arrays.asList(versions), rt -> rt == null ? "null" : rt.toString(table))) + ']',
                                                           sources.contacts(),
                                                           makeResponsesDebugString(partitionKey));
                            throw new AssertionError(details, e);
                        }

                    }

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

            public void close()
            {
                partitionListener.close();
                if (repairedDataTracker != null)
                    repairedDataTracker.verify();
            }
        };
    }
}
