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

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

import com.google.common.collect.Iterables;

import org.assertj.core.api.Assertions;
import org.junit.Test;
import org.mockito.Mockito;
import org.mockito.stubbing.Answer;

import accord.api.RoutingKey;
import accord.local.Command;
import accord.local.Node;
import accord.local.StoreParticipants;
import accord.primitives.Ballot;
import accord.primitives.Deps;
import accord.primitives.FullRoute;
import accord.primitives.KeyDeps;
import accord.primitives.Keys;
import accord.primitives.PartialTxn;
import accord.primitives.RangeDeps;
import accord.primitives.Ranges;
import accord.primitives.Routable;
import accord.primitives.SaveStatus;
import accord.primitives.Status;
import accord.primitives.Timestamp;
import accord.primitives.Txn;
import accord.primitives.TxnId;

import org.apache.cassandra.config.CassandraRelevantProperties;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.cql3.CQLTester;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.db.marshal.Int32Type;
import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.LocalPartitioner;
import org.apache.cassandra.dht.ReversedLongLocalPartitioner;
import org.apache.cassandra.exceptions.InvalidRequestException;
import org.apache.cassandra.schema.MemtableParams;
import org.apache.cassandra.schema.Schema;
import org.apache.cassandra.schema.SchemaProvider;
import org.apache.cassandra.schema.TableId;
import org.apache.cassandra.service.accord.AccordKeyspace.CommandsForKeyAccessor;
import org.apache.cassandra.service.accord.api.TokenKey;
import org.apache.cassandra.utils.CassandraGenerators;

import static accord.local.Command.Committed.committed;
import static accord.utils.Property.qt;
import static org.apache.cassandra.config.DatabaseDescriptor.getPartitioner;
import static org.apache.cassandra.config.TestDatabaseDescriptor.setUnsafeSelectedSSTableFormat;
import static org.apache.cassandra.db.ColumnFamilyStore.FlushReason.UNIT_TESTS;
import static org.apache.cassandra.distributed.test.log.ClusterMetadataTestHelper.setMemtable;
import static org.apache.cassandra.schema.SchemaConstants.ACCORD_KEYSPACE_NAME;
import static org.apache.cassandra.service.accord.AccordKeyspace.CommandsForKeyAccessor.findAllKeysBetween;
import static org.apache.cassandra.service.accord.AccordKeyspace.CommandsForKeyAccessor.makeSystemTableKeyBytes;
import static org.apache.cassandra.service.accord.AccordTestUtils.createTxn;
import static org.apache.cassandra.utils.AbstractTypeGenerators.getTypeSupport;
import static org.apache.cassandra.utils.AccordGenerators.fromQT;

public class AccordKeyspaceTest extends CQLTester.InMemory
{
    static
    {
        // since this test does frequent truncates, the info table gets updated and forced flushed... which is 90% of the cost of this test...
        // this flag disables that flush
        CassandraRelevantProperties.UNSAFE_SYSTEM.setBoolean(true);
    }

    @Test
    public void serde()
    {
        AtomicLong now = new AtomicLong();

        String tableName = createTable("CREATE TABLE %s (k int, c int, v int, PRIMARY KEY (k, c)) WITH transactional_mode = 'full'");
        TableId tableId = Schema.instance.getTableMetadata(KEYSPACE, tableName).id;
        Ranges scope = Ranges.of(TokenRange.create(TokenKey.min(tableId, getPartitioner()), TokenKey.max(tableId, getPartitioner())));

        AccordCommandStore store = AccordTestUtils.createAccordCommandStore(now::incrementAndGet, KEYSPACE, tableName);

        TxnId id = new TxnId(Timestamp.fromValues(1, 42, new Node.Id(1)), Txn.Kind.Read, Routable.Domain.Key);

        Txn txn = createTxn(wrapInTxn(String.format("SELECT * FROM %s.%s WHERE k=? LIMIT 1", KEYSPACE, tableName)), Collections.singletonList(42));

        PartialTxn partialTxn = txn.slice(scope, true);
        RoutingKey routingKey = partialTxn.keys().get(0).asKey().toUnseekable();
        FullRoute<?> route = partialTxn.keys().toRoute(routingKey);
        StoreParticipants participants = StoreParticipants.all(route);
        Deps deps = new Deps(KeyDeps.none(((Keys) txn.keys()).toParticipants()), RangeDeps.NONE);

        Command.WaitingOn waitingOn = null;

        Command.Committed committed = committed(id, SaveStatus.Committed, Status.Durability.NotDurable,
                                                participants, Ballot.ZERO, id, partialTxn, deps.intersecting(scope),
                                                Ballot.ZERO, waitingOn);
        AccordSafeCommand safeCommand = new AccordSafeCommand(AccordTestUtils.loaded(id, null));
        safeCommand.set(committed);

        AccordTestUtils.appendCommandsBlocking(store, null, committed);

        Command loaded = store.loadCommand(id);
        Assertions.assertThat(loaded).isEqualTo(committed);
    }

    @Test
    public void findOverlappingKeys()
    {
        var tableIdGen = fromQT(CassandraGenerators.TABLE_ID_GEN);
        var partitionGen = fromQT(CassandraGenerators.partitioners().assuming(IPartitioner::accordSupported));

        var sstableFormats = DatabaseDescriptor.getSSTableFormats();
        List<String> sstableFormatNames = new ArrayList<>(sstableFormats.keySet());
        sstableFormatNames.sort(Comparator.naturalOrder());

        List<String> memtableFormats = MemtableParams.knownDefinitions().stream()
                                                     .filter(name -> !name.startsWith("test_") && !name.equals("default"))
                                                     .sorted()
                                                     .collect(Collectors.toList());

        qt().check(rs -> {
            AccordKeyspace.unsafeClear();
            // control SSTable format
            setUnsafeSelectedSSTableFormat(sstableFormats.get(rs.pick(sstableFormatNames)));
            // control memtable format
            setMemtable(ACCORD_KEYSPACE_NAME, "commands_for_key", rs.pick(memtableFormats));

            // define the tables w/ partitioners for the test
            // this uses the ability to override the SchemaProvider for the keyspace and only defines the single API call expected: getTablePartitioner
            TreeMap<TableId, IPartitioner> tables = new TreeMap<>();
            int numStores = rs.nextInt(1, 3);
            int numTables = numStores == 1 ? 1 : rs.nextInt(1, numStores);
            for (int i = 0; i < numTables; i++)
            {
                var tableId = tableIdGen.next(rs);
                while (tables.containsKey(tableId))
                    tableId = tableIdGen.next(rs);
                tables.put(tableId, partitionGen.next(rs));
            }

            TreeMap<Integer, TableId> storeTableIds = new TreeMap<>();
            for (int i = 0; i < numStores; i++)
            {
                int tableIdx = rs.nextInt(tables.size());
                TableId tableId = Iterables.get(tables.keySet(), tableIdx);
                storeTableIds.put(i, tableId);
            }
            SchemaProvider schema = Mockito.mock(SchemaProvider.class);
            Mockito.when(schema.getTablePartitioner(Mockito.any())).thenAnswer((Answer<IPartitioner>) invocationOnMock -> tables.get(invocationOnMock.getArgument(0)));

            // The model of the DB
            TreeMap<Integer, SortedSet<TokenKey>> storesToKeys = new TreeMap<>();
            // write to the table and the model
            for (int i = 0, numKeys = rs.nextInt(10, 20); i < numKeys; i++)
            {
                int store = rs.nextInt(0, numStores);
                var keys = storesToKeys.computeIfAbsent(store, ignore -> new TreeSet<>());
                TokenKey pk = null;
                // LocalPartitioner may have a type with a very small domain (boolean, vector<boolean, 1>, etc.), so need to bound the attempts
                // else this will loop forever...
                for (int attempt = 0; attempt < 10; attempt++)
                {
                    TableId tableId = storeTableIds.get(store);
                    IPartitioner partitioner = tables.get(tableId);
                    ByteBuffer data;
                    if (partitioner instanceof ReversedLongLocalPartitioner)
                        data = fromQT(CassandraGenerators.reversedLongLocalKeys()).next(rs);
                    else if (partitioner instanceof LocalPartitioner)
                        data = fromQT(getTypeSupport(partitioner.getTokenValidator()).bytesGen()).next(rs);
                    else
                        data = Int32Type.instance.decompose(rs.nextInt());
                    TokenKey key = new TokenKey(tableId, tables.get(tableId).decorateKey(data).getToken());
                    if (keys.add(key))
                    {
                        pk = key;
                        break;
                    }
                }
                if (pk != null)
                {
                    try
                    {
                        // using Mutation directly (what we do in Accord) can break when user data is too large; leading to data loss
                        // The memtable will allow the write, but it will be dropped when writing to the SSTable...
                        execute("INSERT INTO system_accord.commands_for_key (key) VALUES (?)",
                                makeSystemTableKeyBytes(store, pk));
                    }
                    catch (IllegalArgumentException | InvalidRequestException e)
                    {
                        // Sometimes the types are too large (LocalPartitioner) so the mutation gets rejected... just ignore those cases
                        // Length 69912 > max length 65535
                        String msg = e.getMessage();
                        if (msg != null)
                        {
                            if ((msg.startsWith("Length ") && msg.endsWith("> max length 65535")) // Clustering was rejected
                                || (msg.startsWith("Key length of ") && msg.endsWith(" is longer than maximum of 65535"))) // Partition was rejected
                            {
                                // failed to add
                                keys.remove(pk);
                                continue;
                            }
                        }
                        throw e;
                    }
                }
            }

            // read from the table and validate it matches the model
            for (int read = 0; read < 2; read++) // read=0 is memtable, read=1 is sstable
            {
                {
                    // Make sure no data was lost
                    // An issue was found that system mutations bypass checks so make their way to the Memtable, but when we flush to SSTable
                    // they get filtered out, causing data loss... This check is here to make sure that the data is present (test covers Memtable + SStable)
                    // in the storage before checking if the filtering logic is correct
                    TreeMap<Integer, SortedSet<ByteBuffer>> expectedCqlStoresToKeys = new TreeMap<>();
                    for (var e : storesToKeys.entrySet())
                    {
                        int store = e.getKey();
                        SortedSet<TokenKey> keys = e.getValue();
                        if (keys.isEmpty())
                            continue;
                        expectedCqlStoresToKeys.put(store, new TreeSet<>(keys.stream().map(key -> makeSystemTableKeyBytes(store, key)).collect(Collectors.toList())));
                    }

                    // make sure no data loss... when this test was written sstable had all the rows but the sstable didn't... this
                    // is mostly a santity check to detect that case early
                    var resultSet = execute("SELECT key FROM system_accord.commands_for_key ALLOW FILTERING");
                    TreeMap<Integer, SortedSet<ByteBuffer>> cqlStoresToKeys = new TreeMap<>();
                    for (var row : resultSet)
                    {
                        ByteBuffer bb = row.getBytes("key");
                        int storeId = CommandsForKeyAccessor.getCommandStoreId(bb);
                        cqlStoresToKeys.computeIfAbsent(storeId, ignore -> new TreeSet<>()).add(bb);
                    }
                    Assertions.assertThat(cqlStoresToKeys).isEqualTo(expectedCqlStoresToKeys);
                }

                for (int i = 0, queries = rs.nextInt(1, 5); i < queries; i++)
                {
                    int store = rs.pickOrderedSet(storesToKeys.navigableKeySet());
                    var keysForStore = new ArrayList<>(storesToKeys.get(store));
                    if (keysForStore.isEmpty())
                        continue;

                    int offset;
                    int offsetEnd;
                    if (keysForStore.size() == 1)
                    {
                        offset = 0;
                        offsetEnd = 1;
                    }
                    else
                    {
                        offset = rs.nextInt(0, keysForStore.size());
                        offsetEnd = rs.nextInt(offset, keysForStore.size()) + 1;
                    }
                    List<TokenKey> expected = keysForStore.subList(offset, offsetEnd);
                    TokenKey start = expected.get(0);
                    TokenKey end = expected.get(expected.size() - 1);

                    List<TokenKey> actual = new ArrayList<>();
                    findAllKeysBetween(store, storeTableIds.get(store), start.token().getPartitioner(), start, true, end, true, actual::add);
                    Assertions.assertThat(actual).isEqualTo(expected);
                }

                if (read == 0)
                    Keyspace.open(ACCORD_KEYSPACE_NAME).getColumnFamilyStore("commands_for_key").forceBlockingFlush(UNIT_TESTS);
            }
        });
    }
}