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

import java.util.ArrayList;
import java.util.List;

import com.google.common.annotations.VisibleForTesting;

import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.util.FileUtils;

/**
 * {@link KeyRangeConcatIterator} takes a list of sorted range iterators and concatenates them, leaving duplicates in
 * place, to produce a new stably sorted iterator. Duplicates are eliminated later in
 * {@link org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher}
 * as results from multiple SSTable indexes and their respective segments are consumed.
 * <p>
 * ex. (1, 2, 3) + (3, 3, 4, 5) -> (1, 2, 3, 3, 3, 4, 5)
 * ex. (1, 2, 2, 3) + (3, 4, 4, 6, 6, 7) -> (1, 2, 2, 3, 3, 4, 4, 6, 6, 7)
 */
public class KeyRangeConcatIterator extends KeyRangeIterator
{
    public static final String MUST_BE_SORTED_ERROR = "RangeIterator must be sorted, previous max: %s, next min: %s";
    private final List<KeyRangeIterator> ranges;

    private int current;

    protected KeyRangeConcatIterator(KeyRangeIterator.Builder.Statistics statistics, List<KeyRangeIterator> ranges, Runnable onClose)
    {
        super(statistics, onClose);

        if (ranges.isEmpty())
            throw new IllegalArgumentException("Cannot concatenate empty list of ranges");

        this.current = 0;
        this.ranges = ranges;
    }

    @Override
    protected void performSkipTo(PrimaryKey nextKey)
    {
        while (current < ranges.size())
        {
            KeyRangeIterator currentIterator = ranges.get(current);

            if (currentIterator.hasNext() && currentIterator.peek().compareTo(nextKey, false) >= 0)
                break;

            if (currentIterator.getMaximum().compareTo(nextKey, false) >= 0)
            {
                currentIterator.skipTo(nextKey);
                break;
            }

            current++;
        }
    }

    @Override
    protected PrimaryKey computeNext()
    {
        while (current < ranges.size())
        {
            KeyRangeIterator currentIterator = ranges.get(current);

            if (currentIterator.hasNext())
                return currentIterator.next();

            current++;
        }

        return endOfData();
    }

    @Override
    public void close()
    {
        super.close();

        // due to lazy key fetching, we cannot close iterator immediately
        FileUtils.closeQuietly(ranges);
    }

    public static Builder builder(int size)
    {
        return builder(size, () -> {});
    }

    public static Builder builder(int size, Runnable onClose)
    {
        return new Builder(size, onClose);
    }

    @VisibleForTesting
    public static class Builder extends KeyRangeIterator.Builder
    {
        // We can use a list because the iterators are already in order
        private final List<KeyRangeIterator> ranges;

        Builder(int size, Runnable onClose)
        {
            super(new ConcatStatistics(), onClose);
            this.ranges = new ArrayList<>(size);
        }

        @Override
        public KeyRangeIterator.Builder add(KeyRangeIterator range)
        {
            if (range == null)
                return this;

            if (range.getMaxKeys() > 0)
                ranges.add(range);
            else
                FileUtils.closeQuietly(range);
            statistics.update(range);

            return this;
        }

        @Override
        public int rangeCount()
        {
            return ranges.size();
        }

        @Override
        public void cleanup()
        {
            super.cleanup();
            FileUtils.closeQuietly(ranges);
        }

        @Override
        protected KeyRangeIterator buildIterator()
        {
            if (rangeCount() == 0)
            {
                onClose.run();
                return empty();
            }
            if (rangeCount() == 1)
            {
                KeyRangeIterator single = ranges.get(0);
                single.setOnClose(onClose);
                return single;
            }

            return new KeyRangeConcatIterator(statistics, ranges, onClose);
        }
    }

    private static class ConcatStatistics extends KeyRangeIterator.Builder.Statistics
    {
        @Override
        public void update(KeyRangeIterator range)
        {
            // range iterators should be sorted, but previous max must not be greater than next min.
            if (range.getMaxKeys() > 0)
            {
                if (count == 0)
                {
                    min = range.getMinimum();
                }
                else if (count > 0 && max.compareTo(range.getMinimum(), false) > 0)
                {
                    throw new IllegalArgumentException(String.format(MUST_BE_SORTED_ERROR, max, range.getMinimum()));
                }

                max = range.getMaximum();
                count += range.getMaxKeys();
            }
        }
    }
}
