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

import java.io.Closeable;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.EnumMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.UUID;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import com.google.common.collect.ImmutableMap;

import org.assertj.core.api.Assertions;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import accord.primitives.TxnId;
import accord.utils.Gen;
import accord.utils.Gens;
import accord.utils.RandomSource;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.ClusteringComparator;
import org.apache.cassandra.db.marshal.ByteBufferAccessor;
import org.apache.cassandra.db.marshal.Int32Type;
import org.apache.cassandra.dht.Murmur3Partitioner;
import org.apache.cassandra.index.accord.CheckpointIntervalArrayIndex.Interval;
import org.apache.cassandra.index.accord.IndexDescriptor.IndexComponent;
import org.apache.cassandra.io.sstable.Descriptor;
import org.apache.cassandra.io.sstable.SequenceBasedSSTableId;
import org.apache.cassandra.io.util.File;
import org.apache.cassandra.io.util.FileHandle;
import org.apache.cassandra.schema.TableId;
import org.apache.cassandra.utils.ByteArrayUtil;
import org.apache.cassandra.utils.bytecomparable.ByteComparable;
import org.apache.cassandra.utils.bytecomparable.ByteSource;
import org.apache.cassandra.utils.bytecomparable.ByteSourceInverse;

import static accord.utils.Property.qt;
import static org.apache.cassandra.utils.ByteBufferUtil.bytes;

public class CheckpointIntervalArrayIndexTest
{
    private static final Logger logger = LoggerFactory.getLogger(CheckpointIntervalArrayIndexTest.class);

    static
    {
        DatabaseDescriptor.toolInitialization();
    }

    private static final byte[] EMPTY = new byte[0];
    private static final TreeSet<Interval> EMPTY_TREE_SET = new TreeSet<>();

    private static final Gen.IntGen MAX_TOKEN_GEN = Gens.pickInt(1 << 14,
                                                                 1 << 16,
                                                                 1 << 20,
                                                                 1 << 30);

    private enum Pattern { RANDOM, NO_OVERLAP, PARTIAL_OVERLAP }

    private static final Gen<Pattern> PATTERN_GEN = Gens.enums().all(Pattern.class);

    @Rule
    public TemporaryFolder folder = new TemporaryFolder();
    private int generation = 0;

    @Test
    public void simple() throws IOException
    {
        int bytesPerKey = Integer.BYTES;
        int bytesPerValue = 0;

        List<Interval> list = new ArrayList<>(10);
        list.add(new Interval(bytes(Integer.MIN_VALUE).array(), bytes(Integer.MAX_VALUE).array(), EMPTY));
        for (int i = 0; i < 10; i++)
            list.add(new Interval(bytes(i).array(), bytes(i + 1).array(), EMPTY));

        try (var searcher = index(bytesPerKey, bytesPerValue, list))
        {
            Set<List<Integer>> expected = Set.of(List.of(-2147483648, 2147483647),
                                                 List.of(2, 3),
                                                 List.of(3, 4));
            Set<List<Integer>> actual = new HashSet<>();
            var stats = searcher.intersects(bytes(2).array(), bytes(4).array(), value -> actual.add(List.of(ByteBuffer.wrap(value.start).getInt(), ByteBuffer.wrap(value.end).getInt())));
            logger.info("Stats: {}", stats);
            Assertions.assertThat(actual).isEqualTo(expected);
        }
    }

    @Test
    public void fuzzSmall()
    {
        var minToken = 0;
        int numRecords = 10;
        qt().withTimeout(Duration.ofSeconds(60)).check(rs -> fuzz(rs, minToken, MAX_TOKEN_GEN.nextInt(rs), PATTERN_GEN.next(rs), numRecords));
    }

    @Test
    public void fuzzMedium()
    {
        var minToken = 0;
        int numRecords = 1_000;
        qt().withTimeout(Duration.ofSeconds(60)).check(rs -> fuzz(rs, minToken, MAX_TOKEN_GEN.nextInt(rs), PATTERN_GEN.next(rs), numRecords));
    }

    private void fuzz(RandomSource rs, int minToken, int maxToken, Pattern pattern, int numRecords) throws IOException
    {
        List<DetailedInterval> intervals = buildIntervals(rs, minToken, maxToken, pattern, numRecords);
        List<? extends Interval> nonContainedRanges = findMissingRanges(intervals);

        try (var searcher = index(Integer.BYTES, Integer.BYTES, intervals))
        {
            for (int i = 0, samples = rs.nextInt(Math.min(10, numRecords), Math.min(10, numRecords) * 10); i < samples; i++)
            {
                SearchContext ctx = rs.decide(.2) ? miss(rs, nonContainedRanges)
                                                  : hit(rs, intervals, pattern);
                Set<Interval> actual = new TreeSet<>();
                try
                {
                    var stats = searcher.intersects(ctx.start, ctx.end, interval -> actual.add(new DetailedInterval(interval)));
                    logger.info("[Pattern={}, size={}, expectedMatches={}, query=[{}, {})] Stats: {}", pattern, intervals.size(), ctx.expected.size(), ctx.a, ctx.b, stats);
                }
                catch (Throwable t)
                {
                    throw new AssertionError(String.format("Failure searching for [%d, %d) from %s", ctx.a, ctx.b, intervals), t);
                }
                Assertions.assertThat(actual).describedAs("search(%d, %d) from %s", ctx.a, ctx.b, intervals).isEqualTo(ctx.expected);
            }
        }
    }

    /**
     * mutable/shared ctx to avoid allocating in a loop...
     */
    private SearchContext searchContext = new SearchContext();

    private SearchContext miss(RandomSource rs, List<? extends Interval> nonContainedRanges)
    {
        var range = rs.pick(nonContainedRanges);
        var s = unbc(range.start);
        var e = unbc(range.end);
        int domain = e - s;
        int a, b;
        if (domain == 1)
        {
            // you can not find multiple values within this range!
            a = s;
            b = e;
        }
        else
        {
            a = e == Integer.MAX_VALUE ? rs.nextInt(s, e) : rs.nextInt(s, e) + 1;
            b = e == Integer.MAX_VALUE ? rs.nextInt(s, e) : rs.nextInt(s, e) + 1;
            for (int i = 0; i < 42 && a == b; i++)
                b = e == Integer.MAX_VALUE ? rs.nextInt(s, e) : rs.nextInt(s, e) + 1;
            if (a == b)
                throw new IllegalStateException("Unable to create missing range: " + range);
            if (b < a)
            {
                var tmp = a;
                a = b;
                b = tmp;
            }
        }
        searchContext.a = a;
        searchContext.b = b;
        searchContext.start = bc(a);
        searchContext.end = bc(b);
        searchContext.expected = EMPTY_TREE_SET;
        return searchContext;
    }

    private SearchContext hit(RandomSource rs, List<DetailedInterval> intervals, Pattern pattern)
    {
        int numRecords =  intervals.size();
        DetailedInterval first, second;
        do
        {
            var offset = rs.nextInt(0, numRecords);
            int endOffset;
            switch (pattern)
            {
                case PARTIAL_OVERLAP:
                case RANDOM:
                    endOffset = offset;
                    break;
                case NO_OVERLAP:
                    endOffset = offset == numRecords - 1 ? offset : offset + rs.nextInt(1, Math.min(3, numRecords - offset));
                    break;
                default:
                    throw new IllegalArgumentException("Unknown pattern: " + pattern);
            }
            first = intervals.get(offset);
            second = intervals.get(endOffset);
        }
        while (first.compareTo(second) == 0 && first.size() == 1);
        int a, b;
        a = rs.nextInt(unbc(first.start), unbc(first.end)) + 1;
        b = rs.nextInt(unbc(second.start), unbc(second.end)) + 1;
        while (a == b)
            b = rs.nextInt(unbc(second.start), unbc(second.end)) + 1;
        if (b < a)
        {
            var tmp = b;
            b = a;
            a = tmp;
        }

        searchContext.start = bc(a);
        searchContext.end = bc(b);

        searchContext.expected = intervals.stream().filter(i -> i.intersects(searchContext.start, searchContext.end)).collect(Collectors.toCollection(TreeSet::new));
        Assertions.assertThat(searchContext.expected).isNotEmpty();
        return searchContext;
    }

    private static List<DetailedInterval> buildIntervals(RandomSource rs, int minToken, int maxToken, Pattern pattern, int numRecords)
    {
        List<DetailedInterval> intervals = new ArrayList<>(numRecords);
        {
            var domain = maxToken - minToken + 1;
            var delta = domain / numRecords;
            var sub_delta = delta / 2;
            for (int i = 0; i < numRecords; i++)
            {
                switch (pattern)
                {
                    case RANDOM:
                    {
                        var start = rs.nextInt(minToken, maxToken);
                        var remaining = maxToken - start;
                        var end = start + (remaining == 1 ? 1 : rs.nextInt(1, remaining));
                        intervals.add(new DetailedInterval(bc(start), bc(end), bytes(i).array()));
                    }
                    break;
                    case NO_OVERLAP:
                    {
                        var start = delta * i;
                        var end = start + sub_delta;
                        intervals.add(new DetailedInterval(bc(start), bc(end), bytes(i).array()));
                    }
                    break;
                    case PARTIAL_OVERLAP:
                    {
                        if (i > 1 && rs.decide(.2))
                        {
                            // overlap
                            DetailedInterval start, end;
                            do
                            {
                                int numOverlaps = rs.nextInt(1, Math.min(3, intervals.size()));
                                int offset = rs.nextInt(0, intervals.size() - numOverlaps);
                                start = intervals.get(offset);
                                end = intervals.get(offset + numOverlaps);
                            }
                            while (start.compareStart(end) == 0 && start.size() == 1);
                            var a = rs.nextInt(unbc(start.start), unbc(start.end)) + 1;
                            var b = rs.nextInt(unbc(end.start), unbc(end.end)) + 1;
                            if (a == b && end.size() == 1)
                            {
                                while (a == b)
                                    a = rs.nextInt(unbc(start.start), unbc(start.end)) + 1;
                            }
                            else
                            {
                                while (a == b)
                                    b = rs.nextInt(unbc(end.start), unbc(end.end)) + 1;
                            }
                            if (a > b)
                            {
                                var tmp = a;
                                a = b;
                                b = tmp;
                            }
                            intervals.add(new DetailedInterval(bc(a), bc(b), bytes(i).array()));
                            intervals.sort(Comparator.naturalOrder()); // so partial can work next time
                        }
                        else
                        {
                            // no overlap
                            var start = delta * i;
                            var end = start + sub_delta;
                            intervals.add(new DetailedInterval(bc(start), bc(end), bytes(i).array()));
                        }
                    }
                    break;
                    default:
                        throw new IllegalArgumentException("Unknown pattern: " + pattern);
                }
            }
            intervals.sort(Comparator.naturalOrder());
        }
        return intervals;
    }

    private static List<DetailedInterval> findMissingRanges(List<? extends Interval> intervals)
    {
        List<DetailedInterval> list = new ArrayList<>();
        list.add(new DetailedInterval(bc(Integer.MIN_VALUE), intervals.get(0).start, bytes(0).array()));
        // track current visable coverage
        int end = unbc(intervals.get(0).end);
        for (var i : intervals)
        {
            int istar = unbc(i.start);
            int iend = unbc(i.end);
            if (end >= istar)
            {
                // current scope includes this range
                end = Math.max(end, iend);
            }
            else
            {
                // range doesn't intersect, and a new start/end are formed!
                list.add(new DetailedInterval(bc(end), bc(istar), bytes(list.size()).array()));
                end = iend;
            }
        }
        list.add(new DetailedInterval(bc(end), bc(Integer.MAX_VALUE), bytes(list.size()).array()));
        return list;
    }

    private static class DetailedInterval extends Interval
    {
        public DetailedInterval(byte[] start, byte[] end, byte[] value)
        {
            super(start, end, value);
        }

        public DetailedInterval(Interval other)
        {
            super(other);
        }

        public int size()
        {
            return unbc(end) - unbc(start);
        }

        @Override
        public String toString()
        {
            return "[" + unbc(start) + ", " + unbc(end) + ") -> " + ByteBuffer.wrap(value).getInt();
        }
    }

    private static byte[] bc(int value)
    {
        ByteBuffer bb = bytes(value);
        var bs = Int32Type.instance.asComparableBytes(ByteBufferAccessor.instance, bb, ByteComparable.Version.OSS50);
        return ByteSourceInverse.readBytes(bs);
    }

    private static int unbc(byte[] bc)
    {
        return Int32Type.instance.fromComparableBytes(ByteSource.peekable(ByteSource.fixedLength(bc)), ByteComparable.Version.OSS50).getInt();
    }

    @SuppressWarnings({ "IOResourceOpenedButNotSafelyClosed", "resource" })
    private Searcher index(int bytesPerKey, int bytesPerValue, List<? extends Interval> sortedIntervals) throws IOException
    {
        IndexDescriptor descriptor = nextDescriptor();

        var writer = new CheckpointIntervalArrayIndex.SegmentWriter(descriptor, bytesPerKey, bytesPerValue);
        var metas = writer.write(sortedIntervals.toArray(Interval[]::new));

        // going through the RouteIndexFormat isn't required for this test, but it helps improve coverage there...
        Segment segment = new Segment(ImmutableMap.of(new Key(0, TableId.fromUUID(new UUID(0, 0))), new Segment.Metadata(metas, ByteArrayUtil.EMPTY_BYTE_ARRAY, ByteArrayUtil.EMPTY_BYTE_ARRAY, TxnId.NONE, TxnId.MAX, TxnId.NONE)));
        RouteIndexFormat.appendSegment(descriptor, segment);

        Map<IndexComponent, FileHandle> files = new EnumMap<>(IndexComponent.class);
        for (IndexComponent c : descriptor.getLiveComponents())
            files.put(c, new FileHandle.Builder(descriptor.fileFor(c)).mmapped().complete());
        List<Segment> segments = RouteIndexFormat.readSegments(files);
        files.remove(IndexComponent.SEGMENT).close();
        files.remove(IndexComponent.METADATA).close();

        var searcher = new CheckpointIntervalArrayIndex.SegmentSearcher(files.get(IndexComponent.CINTIA_SORTED_LIST).sharedCopy(), metas.get(IndexComponent.CINTIA_SORTED_LIST).offset,
                                                                        files.get(IndexComponent.CINTIA_CHECKPOINTS).sharedCopy(), metas.get(IndexComponent.CINTIA_CHECKPOINTS).offset);
        return new Searcher()
        {
            @Override
            public CheckpointIntervalArrayIndex.Stats intersects(byte[] start, byte[] end, Consumer<Interval> callback) throws IOException
            {
                return searcher.intersects(start, end, callback);
            }

            @Override
            public void close()
            {
                searcher.close();
                for (var fh : files.values())
                    fh.close();
            }
        };
    }

    private IndexDescriptor nextDescriptor()
    {
        return IndexDescriptor.create(new Descriptor(new File(folder.getRoot()), "test", "test", new SequenceBasedSSTableId(generation++)),
                                      Murmur3Partitioner.instance,
                                      new ClusteringComparator());
    }

    private static class SearchContext
    {
        TreeSet<Interval> expected;
        byte[] start, end;
        int a, b;
    }

    public interface Searcher extends Closeable
    {
        CheckpointIntervalArrayIndex.Stats intersects(byte[] start, byte[] end, Consumer<Interval> callback) throws IOException;
    }
}