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

import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.NavigableSet;
import java.util.function.Consumer;

import javax.annotation.Nullable;

import org.quicktheories.generators.SourceDSL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import accord.utils.DefaultRandom;
import accord.utils.Gen;
import accord.utils.Gens;
import accord.utils.RandomSource;

import org.apache.cassandra.cql3.ast.CQLFormatter;
import org.apache.cassandra.cql3.ast.Mutation;
import org.apache.cassandra.cql3.ast.Select;
import org.apache.cassandra.cql3.ast.StandardVisitors;
import org.apache.cassandra.cql3.ast.Statement;
import org.apache.cassandra.cql3.ast.Symbol;
import org.apache.cassandra.cql3.ast.Txn;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.marshal.AsciiType;
import org.apache.cassandra.db.marshal.BytesType;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.dht.Murmur3Partitioner;
import org.apache.cassandra.distributed.Cluster;
import org.apache.cassandra.distributed.api.ConsistencyLevel;
import org.apache.cassandra.distributed.api.IInvokableInstance;
import org.apache.cassandra.distributed.api.IIsolatedExecutor;
import org.apache.cassandra.distributed.impl.Query;
import org.apache.cassandra.distributed.impl.RowUtil;
import org.apache.cassandra.harry.model.ASTSingleTableModel;
import org.apache.cassandra.harry.model.BytesPartitionState;
import org.apache.cassandra.harry.util.StringUtils;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.service.consensus.TransactionalMode;
import org.apache.cassandra.service.reads.repair.ReadRepairStrategy;
import org.apache.cassandra.simulator.AbstractSimulation;
import org.apache.cassandra.simulator.Action;
import org.apache.cassandra.simulator.ActionList;
import org.apache.cassandra.simulator.ActionSchedule;
import org.apache.cassandra.simulator.Actions;
import org.apache.cassandra.simulator.RunnableActionScheduler;
import org.apache.cassandra.simulator.cluster.ClusterActions;
import org.apache.cassandra.simulator.systems.SimulatedActionCallable;
import org.apache.cassandra.simulator.systems.SimulatedSystems;
import org.apache.cassandra.utils.ASTGenerators;
import org.apache.cassandra.utils.AbstractTypeGenerators;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.CassandraGenerators;
import org.apache.cassandra.utils.FastByteOperations;
import org.apache.cassandra.utils.Generators;

import static org.apache.cassandra.simulator.cluster.ClusterActions.InitialConfiguration.initializeAll;
import static org.apache.cassandra.utils.AbstractTypeGenerators.overridePrimitiveTypeSupport;
import static org.apache.cassandra.utils.AbstractTypeGenerators.stringComparator;
import static org.apache.cassandra.utils.Generators.toGen;

public class SingleTableASTSimulation extends SimulationTestBase.SimpleSimulation
{
    private static final int MAX_STEPS = 5000;

    static
    {
        // limit text/bytes so they are not too big; mostly for debugging than anything
        overridePrimitiveTypeSupport(AsciiType.instance, AbstractTypeGenerators.TypeSupport.of(AsciiType.instance, SourceDSL.strings().ascii().ofLengthBetween(1, 10), stringComparator(AsciiType.instance)));
        overridePrimitiveTypeSupport(UTF8Type.instance, AbstractTypeGenerators.TypeSupport.of(UTF8Type.instance, Generators.utf8(1, 10), stringComparator(UTF8Type.instance)));
        overridePrimitiveTypeSupport(BytesType.instance, AbstractTypeGenerators.TypeSupport.of(BytesType.instance, Generators.bytes(1, 10), FastByteOperations::compareUnsigned));
    }

    private final RandomSource rs;
    private ASTRunner runner;

    protected SingleTableASTSimulation(SimulatedSystems simulated, RunnableActionScheduler scheduler, Cluster cluster, ClusterActions.Options options)
    {
        super(simulated, scheduler, cluster, options);
        this.rs = new DefaultRandom(simulated.random.uniform(Long.MIN_VALUE, Long.MAX_VALUE));

        cluster.stream().forEach((IInvokableInstance i) -> simulated.failureDetector.markUp(i.config().broadcastAddress()));
    }

    @Override
    protected ActionSchedule.Mode mode()
    {
        return ActionSchedule.Mode.STREAM_LIMITED;
    }

    @Override
    protected ActionList initialize()
    {
        List<Action> actions = new ArrayList<>();
        actions.add(clusterActions.initializeCluster(initializeAll(cluster.size())));

        int[] dcSizes = new int[clusterActions.snitch.dcCount()];
        for (int nodeId = 1; nodeId <= cluster.size(); nodeId++)
            dcSizes[clusterActions.snitch.dcOf(nodeId)]++;

        actions.addAll(setupTable(dcSizes));
        return ActionList.of(actions);
    }

    protected List<Action> setupTable(int[] dcSizes)
    {
        List<Action> actions = new ArrayList<>();

        StringBuilder createKeyspace = new StringBuilder("CREATE KEYSPACE ks WITH replication = {'class': 'NetworkTopologyStrategy'"); // when the simulation starts the RF gets populated
        for (int i = 0; i < dcSizes.length; i++)
        {
            String name = clusterActions.snitch.nameOfDc(i);
            createKeyspace.append(", '").append(name).append("': ").append(Math.min(3, dcSizes[i]));
        }
        createKeyspace.append("};");
        actions.add(clusterActions.schemaChange(1, createKeyspace.toString()));

        TableMetadata metadata = defineTable(rs, "ks");

        CassandraGenerators.visitUDTs(metadata, udt -> actions.add(clusterActions.schemaChange(1, udt.toCqlString(false, false, false))));
        actions.add(clusterActions.schemaChange(1, metadata.toCqlString(false, false, false)));

        this.runner = new ASTRunner(metadata, rs, MAX_STEPS, this);

        return actions;
    }

    protected AbstractTypeGenerators.TypeGenBuilder supportedTypes()
    {
        return AbstractTypeGenerators.withoutUnsafeEquality()
                                     .withTypeKinds(AbstractTypeGenerators.TypeKind.PRIMITIVE);
    }

    protected TableMetadata defineTable(RandomSource rs, String ks)
    {
        TableMetadata tbl = toGen(new CassandraGenerators.TableMetadataBuilder()
                                  .withTableKinds(TableMetadata.Kind.REGULAR)
                                  .withKnownMemtables()
                                  .withKeyspaceName(ks).withTableName("tbl")
                                  .withSimpleColumnNames()
                                  .withDefaultTypeGen(supportedTypes())
                                  .withPartitioner(Murmur3Partitioner.instance)
                                  .build())
                            .next(rs);
        return tbl.unbuild().params(tbl.params.unbuild().readRepair(ReadRepairStrategy.NONE).build()).build();
    }

    @Override
    protected ActionList execute()
    {
        return ActionList.of(test());
    }

    protected Action test()
    {
        return Actions.stream(1, runner::next);
    }

    public static class FullAccordSingleTableASTSimulation extends SingleTableASTSimulation
    {
        public FullAccordSingleTableASTSimulation(SimulatedSystems simulated, RunnableActionScheduler scheduler, Cluster cluster, ClusterActions.Options options)
        {
            super(simulated, scheduler, cluster, options);
        }

        @Override
        protected TableMetadata defineTable(RandomSource rs, String ks)
        {
            TableMetadata metadata = super.defineTable(rs, ks);
            return metadata.unbuild()
                           .params(metadata.params.unbuild()
                                                  .transactionalMode(TransactionalMode.full)
                                                  .build())
                           .build();
        }
    }

    public static class MixedReadsAccordSingleTableASTSimulation extends SingleTableASTSimulation
    {
        public MixedReadsAccordSingleTableASTSimulation(SimulatedSystems simulated, RunnableActionScheduler scheduler, Cluster cluster, ClusterActions.Options options)
        {
            super(simulated, scheduler, cluster, options);
        }

        @Override
        protected TableMetadata defineTable(RandomSource rs, String ks)
        {
            TableMetadata metadata = super.defineTable(rs, ks);
            return metadata.unbuild()
                           .params(metadata.params.unbuild()
                                                  .transactionalMode(TransactionalMode.mixed_reads)
                                                  .build())
                           .build();
        }
    }

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

        private final TableMetadata metadata;
        private final ASTSingleTableModel model;
        private final RandomSource rs;
        private final int maxSteps;
        private final Gen<Action> commands;
        private final AbstractSimulation simulation;
        private int step;

        public ASTRunner(TableMetadata metadata, RandomSource rs, int maxSteps, AbstractSimulation simulation)
        {
            this.metadata = metadata;
            this.model = new ASTSingleTableModel(metadata);
            this.rs = rs;
            this.maxSteps = maxSteps;
            this.simulation = simulation;

            Gen.IntGen nodeGen = r -> r.nextInt(0, simulation.cluster.size()) + 1;

            List<LinkedHashMap<Symbol, Object>> uniquePartitions = Gens.lists(toGen(ASTGenerators.columnValues(model.factory.partitionColumns)))
                                                                       .uniqueBestEffort()
                                                                       .ofSize(rs.nextInt(1, 20))
                                                                       .next(rs);
            Gen<Action> mutationGen = toGen(ASTGenerators.mutationBuilder(rs, model, uniquePartitions, i -> null)
                                                         .disallowListElementAccessForUpdateSet() //TODO (coverage): CASSANDRA-20828 found an issue with multi cell list type timestamp handling, so make sure accord doesn't hit this
                                                         .build())
                                      .map(mutation -> query(mutation));

            Gen<Action> selectPartitionGen = Gens.pick(uniquePartitions)
                                                 .map(partition -> query(select(partition).build()));
            Gen<Action> selectRowGen = Gens.pick(uniquePartitions).map(this::selectRow);
            Gen<Action> txnGen = Generators.toGen(new ASTGenerators.ModelBasedTxnGenBuilder(rs, model, i -> null, SourceDSL.arbitrary().pick(uniquePartitions)).disallowEmpty().build())
                                 .map(txn -> query(txn));
            Gens.OneOfBuilder<Action> commandsBuilder = Gens.<Action>oneOf()
                                                            .add(mutationGen)
                                                            .add(selectPartitionGen)
                                                            .add(i -> query(Select.builder(metadata).build()))
                                                            .add(r -> {
                                                                int nodeId = nodeGen.nextInt(r);
                                                                logger.warn("[step={}] Scheduling flush on node{}", step, nodeId);
                                                                return simulation.clusterActions.flush(nodeId, metadata.keyspace, metadata.name);
                                                            })
                                                            .add(r -> {
                                                                int nodeId = nodeGen.nextInt(r);
                                                                logger.warn("[step={}] Scheduling compact on node{}", step, nodeId);
                                                                return simulation.clusterActions.compact(nodeId, metadata.keyspace, metadata.name);
                                                            });

            if (!model.factory.clusteringColumns.isEmpty())
                commandsBuilder.add(selectRowGen);
            if (metadata.params.transactionalMode.accordIsEnabled)
                commandsBuilder.add(txnGen);
            this.commands = commandsBuilder.buildWithDynamicWeights().next(rs);
        }

        public Action next()
        {
            if (step + 1 >= maxSteps)
            {
                logger.warn("Max steps reached, exiting...");
                return null;
            }
            step++;
            Action next = commands.next(rs);
            // empty actions implies that the model state doesn't have enough data to process the action, so we want to "skip"
            while (next == null)
                next = commands.next(rs);
            return next;
        }

        private Action selectRow(LinkedHashMap<Symbol, Object> partition)
        {

            var builder = select(partition);
            List<Clustering<ByteBuffer>> partitions = model.partitions(builder.build());
            switch (partitions.size())
            {
                case 0:
                    return null;
                case 1:
                    break;
                default:
                    throw new IllegalStateException("Model matched multiple partitions, only 1 is expected");
            }
            BytesPartitionState state = model.get(partitions.get(0));
            if (state == null)
                return null;
            NavigableSet<Clustering<ByteBuffer>> clusteringKeys = state.clusteringKeys();
            if (clusteringKeys.isEmpty())
                return null;
            Clustering<ByteBuffer> clusteringKey = rs.pickOrderedSet(clusteringKeys);
            for (Symbol ck : model.factory.clusteringColumns)
                builder.value(ck, clusteringKey.bufferAt(model.factory.clusteringColumns.indexOf(ck)));
            return query(builder.build());
        }

        private Select.Builder select(LinkedHashMap<Symbol, Object> partition)
        {
            Select.Builder builder = Select.builder().table(metadata);
            for (var e : partition.entrySet())
                builder.value(e.getKey(), e.getValue());
            return builder;
        }

        private Action query(Select select)
        {
            return query(select, ConsistencyLevel.ALL, o -> model.validate(RowUtil.toByteBuffer(o), select));
        }

        private Action query(Mutation mutation)
        {
            return query(mutation, ConsistencyLevel.QUORUM, o -> model.update(mutation));
        }

        private Action query(Txn txn)
        {
            return query(txn, ConsistencyLevel.ALL, o -> model.updateAndValidate(RowUtil.toByteBuffer(o), txn));
        }

        private String humanReadable(Statement stmt, @Nullable String annotate)
        {
            // With UTF-8 some chars can cause printing issues leading to error messages that don't reproduce the original issue.
            // To avoid this problem, always escape the CQL so nothing gets lost
            String cql = StringUtils.escapeControlChars(stmt.visit(StandardVisitors.DEBUG).toCQL(CQLFormatter.None.instance));
            if (annotate != null)
                cql += " -- " + annotate;
            return cql;
        }

        private Action query(Statement statement, ConsistencyLevel cl, Consumer<Object[][]> onSuccess)
        {
            // Always use node 1 since we're in single node mode
            int nodeId = 1;
            String postfix = "on node" + nodeId;
            switch (statement.kind())
            {
                case SELECT: break;
                case MUTATION: {
                    Mutation mutation = statement.asMutation();
                    if (mutation.isCas())
                        postfix += ", would apply " + model.shouldApply(mutation);
                } break;
                case TXN: {
                    Txn txn = statement.asTxn();
                    postfix += ", would apply " + model.shouldApply(txn);
                } break;
                default:
                    throw new UnsupportedOperationException(statement.kind().name());
            }
            logger.warn("[step={}] Executing query: {}", step, humanReadable(statement, postfix));
            return new SimulatedActionCallable<>(statement.getClass().getSimpleName(),
                                                 Action.Modifiers.RELIABLE_NO_TIMEOUTS,
                                                 Action.Modifiers.RELIABLE_NO_TIMEOUTS,
                                                 simulation.simulated,
                                                 simulation.cluster.get(nodeId),
                                                 query(statement, cl))
            {
                final int step_id = step;
                final long createdAtNanos = simulation.simulated.time.nanoTime();

                @Override
                public void accept(Object[][] objects, Throwable throwable)
                {
                    if (throwable != null)
                    {
                        logger.error("[step={}] failed", step_id, throwable);
                        simulated.failures.accept(throwable);
                        return;
                    }
                    logger.warn("[step={}] completed after {}", step_id, Duration.ofNanos(simulation.simulated.time.nanoTime() - createdAtNanos));
                    onSuccess.accept(objects);
                }
            };
        }

        private IIsolatedExecutor.SerializableCallable<Object[][]> query(Statement statement, ConsistencyLevel cl)
        {
            String cql = statement.toCQL();
            ByteBuffer[] encoded = statement.bindsEncoded();
            // Simulator doesn't support ByteBuffer (fails to transfer), so need to convert to byte[], which will get converted to ByteBuffer within Cassandra.
            Object[] binds = new Object[encoded.length];
            for (int i = 0; i < encoded.length; i++)
                binds[i] = encoded[i] == null ? null : ByteBufferUtil.getArray(encoded[i]);
            return () -> {
                Query q = new Query(cql, Long.MIN_VALUE, false, cl, null, binds);
                return q.call().toObjectArrays();
            };
        }
    }
}
