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

import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.assertj.core.api.Assertions;
import org.junit.Before;
import org.junit.Test;

import accord.utils.Gen;
import accord.utils.Gens;
import accord.utils.Invariants;
import accord.utils.LazyToString;
import accord.utils.ReflectionUtils;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.marshal.ByteBufferAccessor;
import org.apache.cassandra.dht.ByteOrderedPartitioner;
import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.Murmur3Partitioner;
import org.apache.cassandra.dht.Murmur3Partitioner.LongToken;
import org.apache.cassandra.dht.RandomPartitioner;
import org.apache.cassandra.dht.RandomPartitioner.BigIntegerToken;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.io.Serializers;
import org.apache.cassandra.io.util.DataInputBuffer;
import org.apache.cassandra.io.util.DataOutputBuffer;
import org.apache.cassandra.schema.TableId;
import org.apache.cassandra.service.accord.api.TokenKey;
import org.apache.cassandra.utils.AccordGenerators;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.CassandraGenerators;

import static accord.utils.Property.qt;
import static org.apache.cassandra.service.accord.api.TokenKey.serializer;
import static org.apache.cassandra.utils.AccordGenerators.fromQT;
import static org.apache.cassandra.utils.CassandraGenerators.partitioners;
import static org.apache.cassandra.utils.CassandraGenerators.token;

public class TokenKeyTest
{
    static
    {
        DatabaseDescriptor.clientInitialization();
    }

    @Before
    public void before()
    {
        // AccordRoutingKey$TokenKey reaches into DD to get partitioner, so need to set that up...
        DatabaseDescriptor.setPartitionerUnsafe(Murmur3Partitioner.instance);
    }

    @Test
    public void beforeIsTokenSentinel()
    {
        qt().forAll(simpleTokenKey()).check(tokenKey -> {
            var t = tokenKey.before();
            Assertions.assertThat(t.isTokenSentinel()).isTrue();
            Assertions.assertThat(t.isTableSentinel()).isEqualTo(tokenKey.isTableSentinel());
            Assertions.assertThat(t.isMin()).isEqualTo(tokenKey.isMin());
            Assertions.assertThat(t.isMax()).isEqualTo(tokenKey.isMax());
            Assertions.assertThat(t.isBefore()).isTrue();
            Assertions.assertThat(t.isAfter()).isFalse();

            Assertions.assertThatThrownBy(() -> t.before());
            Assertions.assertThatThrownBy(() -> t.after());

            Assertions.assertThat(tokenKey.compareTo(t)).isGreaterThan(0);
            Assertions.assertThat(t.compareTo(tokenKey)).isLessThan(0);
        });
    }

    @Test
    public void afterIsTokenSentinel()
    {
        qt().forAll(simpleTokenKey()).check(tokenKey -> {
            var t = tokenKey.after();
            Assertions.assertThat(t.isTokenSentinel()).isTrue();
            Assertions.assertThat(t.isTableSentinel()).isEqualTo(tokenKey.isTableSentinel());
            Assertions.assertThat(t.isMin()).isEqualTo(tokenKey.isMin());
            Assertions.assertThat(t.isMax()).isEqualTo(tokenKey.isMax());
            Assertions.assertThat(t.isBefore()).isFalse();
            Assertions.assertThat(t.isAfter()).isTrue();

            Assertions.assertThat(tokenKey.compareTo(t)).isLessThan(0);
            Assertions.assertThat(t.compareTo(tokenKey)).isGreaterThan(0);
        });
    }

    @Test
    public void serdeSimple()
    {
        Gen<TokenKey> tokenKeyGen = AccordGenerators.allowBeforeAndAfter(simpleTokenKey());
        @SuppressWarnings({ "resource", "IOResourceOpenedButNotSafelyClosed" }) DataOutputBuffer output = new DataOutputBuffer();
        qt().forAll(tokenKeyGen).check(expected -> {
            DatabaseDescriptor.setPartitionerUnsafe(expected.token().getPartitioner());
            Serializers.testSerde(output, serializer, expected);
            testSerdePrefix(output, serializer, expected);
        });
    }

    @Test
    public void serde()
    {
        @SuppressWarnings({ "resource", "IOResourceOpenedButNotSafelyClosed" }) DataOutputBuffer output = new DataOutputBuffer();
        qt().forAll(tokenKeyWithBeforeAndAfterGen())
            .check(key -> {
                IPartitioner partitioner = key.token().getPartitioner();
                DatabaseDescriptor.setPartitionerUnsafe(partitioner);

                Serializers.testSerde(output, serializer, key);
                Assertions.assertThat(serializer.deserializeAndConsume(serializer.serialize(key), partitioner)).isEqualTo(key);
                {
                    TokenKey roundTrip = serializer.deserializeWithPrefixAndImpliedLength(key.prefix(), serializer.serializeWithoutPrefixOrLength(key), partitioner);
                    Assertions.assertThat(roundTrip).isEqualTo(key);
                }
                {
                    TokenKey roundTrip = serializer.deserializeWithPrefixAndImpliedLength(key.prefix(), serializer.serializeWithoutPrefixOrLength(key), ByteBufferAccessor.instance, 0, partitioner);
                    Assertions.assertThat(roundTrip).isEqualTo(key);
                }
                {
                    TokenKey roundTrip = serializer.deserializeWithPrefix(key.prefix(), serializer.serializedSizeWithoutPrefixOrLength(key), serializer.serializeWithoutPrefixOrLength(key), partitioner);
                    Assertions.assertThat(roundTrip).isEqualTo(key);
                }
                {
                    TokenKey roundTrip = serializer.deserializeWithPrefix(key.prefix(), serializer.serializedSizeWithoutPrefixOrLength(key), serializer.serializeWithoutPrefixOrLength(key), ByteBufferAccessor.instance, 0, partitioner);
                    Assertions.assertThat(roundTrip).isEqualTo(key);
                }
                output.clear();
                serializer.serialize(key, output);
                try (DataInputBuffer in = new DataInputBuffer(output.toByteArray()))
                {
                    serializer.skip(in, partitioner);
                    Invariants.require(0 == in.available());
                }
            });
    }

    @Test
    public void compare()
    {
        qt().forAll(tokenKeyGen())
            .check(key -> {
                ByteBuffer keyBytes = serializer.serialize(key);
                for (TokenKey test : mutateAfter(key))
                {
                    ByteBuffer testBytes = serializer.serialize(test);
                    Invariants.require(test.compareTo(key) > 0);
                    Invariants.require(ByteBufferUtil.compareUnsigned(testBytes, keyBytes) > 0);
                }
                for (TokenKey test : mutateBefore(key))
                {
                    ByteBuffer testBytes = serializer.serialize(test);
                    Invariants.require(test.compareTo(key) < 0);
                    Invariants.require(ByteBufferUtil.compareUnsigned(testBytes, keyBytes) < 0);
                }
            });
    }

    @Test
    public void serdeNoTable()
    {
        @SuppressWarnings({ "resource", "IOResourceOpenedButNotSafelyClosed" }) DataOutputBuffer output = new DataOutputBuffer();
        qt().forAll(tokenKeyGen()).check(key -> {
            AccordGenerators.maybeUpdatePartitioner(key);
            Serializers.testSerde(output, TokenKey.noTableSerializer, key, key.table());
        });
    }

    private static Gen<TokenKey> simpleTokenKey()
    {
        return AccordGenerators.partitioner().flatMap(p -> AccordGenerators.routingKeysGen(p));
    }

    private static void testSerdePrefix(DataOutputBuffer output, TokenKey.Serializer serializer, TokenKey input) throws IOException
    {
        output.clear();
        Object expected = input.prefix();
        long expectedSize = serializer.serializedSizeOfPrefix(expected);
        serializer.serializePrefix(expected, output);
        Assertions.assertThat(output.getLength()).describedAs("The serialized size and bytes written do not match").isEqualTo(expectedSize);
        DataInputBuffer in = new DataInputBuffer(output.unsafeGetBufferAndFlip(), false);
        Object read = serializer.deserializePrefix(in);
        Assertions.assertThat(read)
                  .describedAs("The deserialized output does not match the serialized input; difference %s", new LazyToString(() -> ReflectionUtils.recursiveEquals(read, expected).toString()))
                  .isEqualTo(expected);
    }

    private static Gen<TokenKey> tokenKeyGen()
    {
        return fromQT(partitioners()).filter(IPartitioner::accordSupported)
               .flatMap(partitioner -> routingKeyGen(fromQT(CassandraGenerators.TABLE_ID_GEN), fromQT(token(partitioner)), partitioner));
    }

    private static Gen<TokenKey> tokenKeyWithBeforeAndAfterGen()
    {
        return AccordGenerators.allowBeforeAndAfter(tokenKeyGen());
    }

    private List<TokenKey> mutateAfter(TokenKey mutate)
    {
        List<TokenKey> results = new ArrayList<>();
        if (!mutate.isTableSentinel())
        {
            Token token = mutate.token();
            if (token instanceof ByteOrderedPartitioner.BytesToken)
            {
                byte[] bytes = (byte[]) token.getTokenValue();
                bytes = bytes.clone();
                for (int i = 0 ; i < bytes.length ; ++i)
                {
                    if ((bytes[i] & 0xff) != 0xff)
                    {
                        ++bytes[i];
                        add(results, mutate.withToken(new ByteOrderedPartitioner.BytesToken(bytes.clone())));
                        --bytes[i];
                    }
                }
                add(results, mutate.withToken(new ByteOrderedPartitioner.BytesToken(Arrays.copyOf(bytes, bytes.length + 1))));
            }
            else if (token instanceof LongToken)
            {
                long value = token.getLongValue();
                if (value < Long.MAX_VALUE)
                    add(results, mutate.withToken(new LongToken(value + 1)));
                for (long v = 2L; v >= 0 ; v <<= 1)
                {
                    if ((value & v) == 0)
                        add(results, mutate.withToken(new LongToken(value | v)));
                }
                if (value >= 0)
                {
                    long higher = value;
                    while ((higher <<= 8) > value)
                        add(results, mutate.withToken(new LongToken(higher)));
                }
                else
                {
                    for (int i = 1 ; i < 8 ; ++i)
                        add(results, mutate.withToken(new LongToken(value >> (i * 8))));
                }
            }
            else if (token instanceof BigIntegerToken)
            {
                BigInteger value = (BigInteger) token.getTokenValue();
                if (value.compareTo(RandomPartitioner.MAXIMUM) < 0)
                    add(results, mutate.withToken(new BigIntegerToken(value.add(BigInteger.ONE))));
                for (long v = 1L; v >= 0 ; v <<= 1)
                {
                    BigInteger i = BigInteger.valueOf(v);
                    if (value.and(i).equals(BigInteger.ZERO))
                        add(results, mutate.withToken(new BigIntegerToken(value.or(i))));
                }
                BigInteger higher = value;
                while ((higher = higher.shiftLeft(8)).compareTo(RandomPartitioner.MAXIMUM) <= 0)
                    add(results, mutate.withToken(new BigIntegerToken(higher)));
            }
            else throw new UnsupportedOperationException();
        }
        TableId tableId = mutate.table();
        if (tableId.msb() != Long.MAX_VALUE)
            add(results, mutate.withTable(TableId.fromRaw(tableId.msb() + 1, tableId.lsb())));
        if (tableId.lsb() != Long.MAX_VALUE)
            add(results, mutate.withTable(TableId.fromRaw(tableId.msb(), tableId.lsb() + 1)));
        return results;
    }

    private List<TokenKey> mutateBefore(TokenKey mutate)
    {
        List<TokenKey> results = new ArrayList<>();
        if (!mutate.isTableSentinel())
        {
            Token token = mutate.token();
            if (token instanceof ByteOrderedPartitioner.BytesToken)
            {
                byte[] bytes = (byte[]) token.getTokenValue();
                bytes = bytes.clone();
                for (int i = 0 ; i < bytes.length ; ++i)
                {
                    add(results, mutate.withToken(new ByteOrderedPartitioner.BytesToken(Arrays.copyOf(bytes, i))));
                    if ((bytes[i] & 0xff) != 0)
                    {
                        --bytes[i];
                        add(results, mutate.withToken(new ByteOrderedPartitioner.BytesToken(bytes.clone())));
                        ++bytes[i];
                    }
                }
            }
            else if (token instanceof LongToken)
            {
                long value = token.getLongValue();
                if (value > Long.MIN_VALUE)
                    add(results, mutate.withToken(new LongToken(value - 1)));
                for (long v = 2L; v >= 0 ; v <<= 1)
                {
                    if ((value & v) != 0)
                        add(results, mutate.withToken(new LongToken(value & ~v)));
                }
                if (value >= 0)
                {
                    for (int i = 1 ; i < 8 ; ++i)
                        add(results, mutate.withToken(new LongToken(value >>> (i * 8))));
                }
                else
                {
                    for (int i = 0 ; i < 7 ; ++i)
                    {
                        long next = value & (-1L << (i * 8));
                        if (next != value)
                            add(results, mutate.withToken(new LongToken(next)));
                    }
                }
            }
            else if (token instanceof BigIntegerToken)
            {
                BigInteger value = (BigInteger) token.getTokenValue();
                if (value.compareTo(RandomPartitioner.MINIMUM.getTokenValue()) > 0)
                    add(results, mutate.withToken(new BigIntegerToken(value.subtract(BigInteger.ONE))));
                for (long v = 1L; v >= 0 ; v <<= 1)
                {
                    BigInteger i = BigInteger.valueOf(v);
                    if (!value.and(i).equals(BigInteger.ZERO))
                        add(results, mutate.withToken(new BigIntegerToken(value.andNot(i))));
                }
                for (int i = 1 ; i < 8 ; ++i)
                    add(results, mutate.withToken(new BigIntegerToken(value.shiftRight(i * 16))));
            }
            else throw new UnsupportedOperationException();
        }
        TableId tableId = mutate.table();
        if (tableId.msb() != Long.MIN_VALUE)
            add(results, mutate.withTable(TableId.fromRaw(tableId.msb() - 1, tableId.lsb())));
        if (tableId.lsb() != Long.MIN_VALUE)
            add(results, mutate.withTable(TableId.fromRaw(tableId.msb(), tableId.lsb() -1)));
        return results;
    }

    private static void add(List<TokenKey> to, TokenKey vary)
    {
        to.add(vary);
        if (!vary.isTokenSentinel())
        {
            to.add(vary.before());
            to.add(vary.after());
        }
    }

    private static Gen<TokenKey> routingKeyGen(Gen<TableId> tableIdGen, Gen<Token> tokenGen, IPartitioner partitioner)
    {
        Gen<TokenKey> result = AccordGenerators.routingKeyGen(tableIdGen, Gens.enums().all(AccordGenerators.RoutingKeyKind.class), tokenGen, partitioner);
        if (!(partitioner instanceof ByteOrderedPartitioner))
            return result;
        return result.map((rs, k) -> {
            byte[] bytes = (byte[]) k.token().getTokenValue();
            if (bytes.length >= 3)
            {
                while (rs.nextFloat() < 0.25f)
                {
                    int i = rs.nextInt(bytes.length - 2);
                    bytes[i] = 0;
                    bytes[i + 1] = (byte) rs.nextInt(0, TokenKey.Serializer.ESCAPE_BYTE);
                }
            }

            return k;
        });
    }

}