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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentSkipListSet;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import org.junit.Assert;
import org.junit.Test;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.ConsistencyLevel;
import org.apache.cassandra.distributed.api.IIsolatedExecutor;
import org.apache.cassandra.distributed.test.log.CMSTestBase;
import org.apache.cassandra.exceptions.RequestFailure;
import org.apache.cassandra.harry.gen.EntropySource;
import org.apache.cassandra.harry.gen.Surjections;
import org.apache.cassandra.harry.gen.rng.PCGFastPure;
import org.apache.cassandra.harry.gen.rng.PcgRSUFast;
import org.apache.cassandra.harry.gen.rng.RngUtils;
import org.apache.cassandra.harry.model.TokenPlacementModel;
import org.apache.cassandra.locator.InetAddressAndPort;
import org.apache.cassandra.metrics.TCMMetrics;
import org.apache.cassandra.net.ConnectionType;
import org.apache.cassandra.net.Message;
import org.apache.cassandra.net.MessageDelivery;
import org.apache.cassandra.net.RequestCallback;
import org.apache.cassandra.schema.ReplicationParams;
import org.apache.cassandra.tcm.AtomicLongBackedProcessor;
import org.apache.cassandra.tcm.ClusterMetadata;
import org.apache.cassandra.tcm.ClusterMetadataService;
import org.apache.cassandra.tcm.Epoch;
import org.apache.cassandra.tcm.MultiStepOperation;
import org.apache.cassandra.tcm.membership.Location;
import org.apache.cassandra.tcm.membership.NodeAddresses;
import org.apache.cassandra.tcm.membership.NodeId;
import org.apache.cassandra.tcm.membership.NodeVersion;
import org.apache.cassandra.tcm.transformations.PrepareJoin;
import org.apache.cassandra.tcm.transformations.PrepareLeave;
import org.apache.cassandra.tcm.transformations.Register;
import org.apache.cassandra.tcm.transformations.UnsafeJoin;
import org.apache.cassandra.utils.concurrent.Future;

import static org.junit.Assert.assertEquals;

public class ProgressBarrierTest extends CMSTestBase
{
    static
    {
        DatabaseDescriptor.setRpcTimeout(10);
        DatabaseDescriptor.setProgressBarrierTimeout(1000);
        DatabaseDescriptor.setProgressBarrierBackoff(5);
    }

    @Test
    public void testProgressBarrierWithMergingRanges()
    {
        TokenPlacementModel.ReplicationFactor rf = new TokenPlacementModel.SimpleReplicationFactor(1);
        try (CMSTestBase.CMSSut sut = new CMSTestBase.CMSSut(AtomicLongBackedProcessor::new, false, rf))
        {
            TokenPlacementModel.NodeFactory nodeFactory = TokenPlacementModel.nodeFactory();
            List<TokenPlacementModel.Node> nodes = new ArrayList<>();
            for (int i = 1; i < 6; i++)
            {
                TokenPlacementModel.Node node = nodeFactory.make(i, 1, 1).overrideToken(i*100);
                nodes.add(node);
                sut.service.commit(new Register(new NodeAddresses(node.addr()),
                                                new Location(node.dc(), node.rack()),
                                                NodeVersion.CURRENT));
                sut.service.commit(new UnsafeJoin(node.nodeId(),
                                                  Collections.singleton(node.longToken()),
                                                  ClusterMetadataService.instance().placementProvider()));
            }

            // 6 node cluster, with a single RF1 keyspace
            // node2 owns range (100,200] & node3 (200,300]
            // if node2 leaves node3 will acquire its ranges and will then own (100, 300]
            // no other peers are involved in this trivial operation
            // the progress barrier should be looking for consensus from (nodes2, node3)
            TokenPlacementModel.Node node2 = nodes.get(1);
            TokenPlacementModel.Node node3 = nodes.get(2);
            sut.service.commit(new PrepareLeave(node2.nodeId(),
                                                true,
                                                ClusterMetadataService.instance().placementProvider(),
                                                LeaveStreams.Kind.UNBOOTSTRAP));
            UnbootstrapAndLeave leave = (UnbootstrapAndLeave) sut.service.metadata().inProgressSequences.get(node2.nodeId());

            // Internally affectedRanges::toPeers uses the same logic as
            // the progress barrier does to identify the consensus group
            Set<NodeId> consensusGroup = leave.barrier().affectedRanges.toPeers(ReplicationParams.simple(1),
                                                                                sut.service.metadata().placements(),
                                                                                sut.service.metadata().directory);
            assertEquals(Set.of(node2.nodeId(), node3.nodeId()), consensusGroup);
        }
    }

    @Test
    public void testProgressBarrier() throws Throwable
    {
        EntropySource rng = new PcgRSUFast(1L, 1l);
        Supplier<Boolean> respond = bools().toGenerator().bind(rng);
        Supplier<TokenPlacementModel.ReplicationFactor> rfs = combine(ints(0, 3),
                                                                      ints(1, 5),
                                                                      bools(),
                                                                      ints(1, 5),
                                                                      (Integer dcs, Integer nodesPerDc, Boolean addAlternate, Integer nodesPerDcAlt) -> {
                                                                         if (dcs == 0)
                                                                             return new TokenPlacementModel.SimpleReplicationFactor(nodesPerDc);
                                                                         else if (addAlternate && nodesPerDcAlt.intValue() != nodesPerDc.intValue())
                                                                         {
                                                                             int[] perDc = new int[dcs + 1];
                                                                             Arrays.fill(perDc, nodesPerDc);
                                                                             perDc[perDc.length - 1] = nodesPerDcAlt;
                                                                             return new TokenPlacementModel.NtsReplicationFactor(perDc);
                                                                         }
                                                                         else
                                                                         {
                                                                             return new TokenPlacementModel.NtsReplicationFactor(dcs, nodesPerDc);
                                                                         }
                                                                     })
                                                             .toGenerator()
                                                             .bind(rng);
        Supplier<Integer> nodes = ints(15, 20).toGenerator().bind(rng);
        Supplier<ConsistencyLevel> cls = Surjections.pick(
                                                    ConsistencyLevel.ALL,
                                                    ConsistencyLevel.QUORUM,
                                                    ConsistencyLevel.LOCAL_QUORUM,
                                                    ConsistencyLevel.EACH_QUORUM,
                                                    ConsistencyLevel.ONE).toGenerator()
                                                    .bind(rng);

        TokenPlacementModel.NodeFactory nodeFactory = TokenPlacementModel.nodeFactory();
        for (int run = 0; run < 100; run++)
        {
            TokenPlacementModel.ReplicationFactor rf = rfs.get();
            try (CMSTestBase.CMSSut sut = new CMSTestBase.CMSSut(AtomicLongBackedProcessor::new, false, rf))
            {
                List<TokenPlacementModel.Node> allNodes = new ArrayList<>();
                TokenPlacementModel.Node node = null;
                // + 1 since one of the nodes will not be joined yet by the time we create progress barrier, which will fail
                // a check with ALL.
                int nodesInCluster = Math.max(rf.total(), nodes.get()) + 1;
                for (int i = 1; i <= nodesInCluster; i++)
                {
                    node = nodeFactory.make(i, (i % rf.dcs()) + 1, 1);
                    allNodes.add(node);
                    sut.service.commit(new Register(new NodeAddresses(node.addr()), new Location(node.dc(), node.rack()), NodeVersion.CURRENT));
                    if (i < nodesInCluster)
                        sut.service.commit(new UnsafeJoin(node.nodeId(), Collections.singleton(node.longToken()), ClusterMetadataService.instance().placementProvider()));
                }

                sut.service.commit(new PrepareJoin(node.nodeId(), Collections.singleton(node.longToken()), ClusterMetadataService.instance().placementProvider(), true, false));

                for (int check = 0; check < 10; check++)
                {
                    ClusterMetadata metadata = ClusterMetadata.current();
                    ConsistencyLevel cl = cls.get();

                    Set<InetAddressAndPort> responded = new ConcurrentSkipListSet<>();
                    MessageDelivery delivery = new MessageDelivery()
                    {
                        public <REQ, RSP> void sendWithCallback(Message<REQ> message, InetAddressAndPort to, RequestCallback<RSP> cb)
                        {
                            // assert that it is a replcia
                            if (respond.get())
                            {
                                responded.add(to);
                                cb.onResponse((Message<RSP>) message.responseWith(message.epoch()).withFrom(to));
                            }
                            else
                            {
                                cb.onFailure(message.from(), RequestFailure.TIMEOUT);
                            }
                        }

                        public <REQ> void send(Message<REQ> message, InetAddressAndPort to) {}
                        public <REQ, RSP> void sendWithCallback(Message<REQ> message, InetAddressAndPort to, RequestCallback<RSP> cb, ConnectionType specifyConnection) {}
                        public <REQ, RSP> Future<Message<RSP>> sendWithResult(Message<REQ> message, InetAddressAndPort to) { return null; }
                        @Override
                        public <V> void respond(V response, Message<?> message) {}
                    };
                    ProgressBarrier progressBarrier = ((MultiStepOperation<Epoch>)metadata.inProgressSequences.get(node.nodeId()))
                                                      .advance(metadata.epoch)
                                                      .barrier()
                                                      .withMessagingService(delivery);

                    progressBarrier.await(cl, metadata);

                    String dc = metadata.directory.location(node.nodeId()).datacenter;
                    switch (cl)
                    {
                        case ALL:
                        {
                            Set<InetAddressAndPort> replicas = metadata.lockedRanges.locked.get(LockedRanges.keyFor(metadata.epoch))
                                                                                           .toPeers(rf.asKeyspaceParams().replication, metadata.placements(), metadata.directory)
                                                                                           .stream()
                                                                                           .map(n -> metadata.directory.getNodeAddresses(n).broadcastAddress)
                                                                                           .collect(Collectors.toSet());

                            Set<InetAddressAndPort> collected = responded.stream().filter(replicas::contains).collect(Collectors.toSet());
                            int expected = rf.total();
                            Assert.assertTrue(String.format("Should have collected at least %d nodes but got %d." +
                                                            "\nRF: %s" +
                                                            "\nReplicas: %s" +
                                                            "\nNodes:    %s", expected, collected.size(), rf, replicas, collected),
                                              collected.size() >= expected);

                            break;
                        }
                        case QUORUM:
                        {
                            Set<InetAddressAndPort> replicas = metadata.lockedRanges.locked.get(LockedRanges.keyFor(metadata.epoch))
                                                                                           .toPeers(rf.asKeyspaceParams().replication, metadata.placements(), metadata.directory)
                                                                                           .stream()
                                                                                           .map(n -> metadata.directory.getNodeAddresses(n).broadcastAddress)
                                                                                           .collect(Collectors.toSet());

                            Set<InetAddressAndPort> collected = responded.stream().filter(replicas::contains).collect(Collectors.toSet());
                            int expected = rf.total() / 2 + 1;
                            Assert.assertTrue(String.format("Should have collected at least %d nodes but got %d." +
                                                            "\nRF: %s" +
                                                            "\nReplicas: %s" +
                                                            "\nNodes: %s", expected, collected.size(), rf, replicas, collected),
                                              collected.size() >= expected);

                            break;
                        }
                        case LOCAL_QUORUM:
                        {
                            List<InetAddressAndPort> replicas = new ArrayList<>(metadata.lockedRanges.locked.get(LockedRanges.keyFor(metadata.epoch))
                                                                                                            .toPeers(rf.asKeyspaceParams().replication, metadata.placements(), metadata.directory)
                                                                                                            .stream()
                                                                                                            .filter((n) -> metadata.directory.location(n).datacenter.equals(dc))
                                                                                                            .map(n -> metadata.directory.getNodeAddresses(n).broadcastAddress)
                                                                                                            .collect(Collectors.toSet()));
                            replicas.sort(InetAddressAndPort::compareTo);
                            Set<InetAddressAndPort> collected = responded.stream().filter(replicas::contains).collect(Collectors.toSet());
                            int expected;
                            if (rf instanceof TokenPlacementModel.SimpleReplicationFactor)
                                expected = rf.total() / 2 + 1;
                            else
                                expected = rf.asMap().get(dc).totalCount / 2 + 1;
                            Assert.assertTrue(String.format("Should have collected at least %d nodes but got %d." +
                                                            "\nRF: %s" +
                                                            "\nReplicas: %s" +
                                                            "\nCollected: %s. Responded: %s", expected, collected.size(), rf, replicas, collected, responded),
                                              collected.size() >= expected);

                            break;
                        }
                        case EACH_QUORUM:
                        {
                            Map<String, Integer> byDc = new HashMap<>();
                            metadata.lockedRanges.locked.get(LockedRanges.keyFor(metadata.epoch))
                                                        .toPeers(rf.asKeyspaceParams().replication, metadata.placements(), metadata.directory)
                                                        .forEach(n -> byDc.compute(metadata.directory.location(n).datacenter,
                                                                                   (k, v) -> v == null ? 1 : v + 1));

                            if (rf instanceof TokenPlacementModel.SimpleReplicationFactor)
                            {
                                int actual = byDc.get(dc);
                                int expected = rf.asMap().get(dc).totalCount / 2 + 1;
                                Assert.assertTrue(String.format("Shuold have collected at least %d nodes, but got %d." +
                                                                "\nRF: %s" +
                                                                "\nNodes: %s", expected, byDc.size(), rf, byDc),
                                                  actual >= expected);
                            }
                            else
                            {
                                for (Map.Entry<String, Integer> e : byDc.entrySet())
                                {
                                    int actual = e.getValue();
                                    int expected = rf.asMap().get(e.getKey()).totalCount / 2 + 1;
                                    Assert.assertTrue(String.format("Shuold have collected at least %d nodes, but got %d." +
                                                                    "\nRF: %s" +
                                                                    "\nNodes: %s", expected, byDc.size(), rf, byDc),
                                                      actual >= expected);
                                }
                            }
                            break;
                        }
                        case ONE:
                            Set<InetAddressAndPort> replicas = metadata.lockedRanges.locked.get(LockedRanges.keyFor(metadata.epoch))
                                                                                           .toPeers(rf.asKeyspaceParams().replication, metadata.placements(), metadata.directory)
                                                                                           .stream()
                                                                                           .map(n -> metadata.directory.getNodeAddresses(n).broadcastAddress)
                                                                                           .collect(Collectors.toSet());

                            Assert.assertTrue(String.format("Should have collected at least one of the replicas %s, but got %s." +
                                                            "\nRF: %s.\nNodes: %s.",
                                                            replicas, responded, rf, allNodes),
                                              responded.stream().anyMatch(replicas::contains));
                            break;
                    }
                }

            }
        }
    }

    @Test
    public void testProgressBarrierDegradingConsistency() throws Throwable
    {
        TokenPlacementModel.ReplicationFactor rf = new TokenPlacementModel.NtsReplicationFactor(5, 5);
        TokenPlacementModel.NodeFactory nodeFactory = TokenPlacementModel.nodeFactory();

        DatabaseDescriptor.setProgressBarrierMinConsistencyLevel(ConsistencyLevel.ONE);
        try (CMSTestBase.CMSSut sut = new CMSTestBase.CMSSut(AtomicLongBackedProcessor::new, false, rf))
        {
            TokenPlacementModel.Node node = null;
            for (int i = 1; i <= 4; i++)
            {
                node = nodeFactory.make(i, (i % rf.dcs()) + 1, 1);
                sut.service.commit(new Register(new NodeAddresses(node.addr()), new Location(node.dc(), node.rack()), NodeVersion.CURRENT));
                if (i < 4)
                    sut.service.commit(new UnsafeJoin(node.nodeId(), Collections.singleton(node.longToken()), ClusterMetadataService.instance().placementProvider()));
            }

            sut.service.commit(new PrepareJoin(node.nodeId(), Collections.singleton(node.longToken()), ClusterMetadataService.instance().placementProvider(), true, false));

            Set<InetAddressAndPort> responded = new ConcurrentSkipListSet<>();
            MessageDelivery delivery = new MessageDelivery()
            {
                AtomicInteger counter = new AtomicInteger();
                public <REQ, RSP> void sendWithCallback(Message<REQ> message, InetAddressAndPort to, RequestCallback<RSP> cb)
                {
                    if (counter.getAndIncrement() == 0)
                    {
                        responded.add(to);
                        cb.onResponse((Message<RSP>) message.responseWith(message.epoch()));
                    }
                }

                public <REQ> void send(Message<REQ> message, InetAddressAndPort to) {}
                public <REQ, RSP> void sendWithCallback(Message<REQ> message, InetAddressAndPort to, RequestCallback<RSP> cb, ConnectionType specifyConnection) {}
                public <REQ, RSP> Future<Message<RSP>> sendWithResult(Message<REQ> message, InetAddressAndPort to) { return null; }

                @Override
                public <V> void respond(V response, Message<?> message) {}
            };

            ClusterMetadata metadata = ClusterMetadata.current();
            ProgressBarrier progressBarrier = ((MultiStepOperation<Epoch>) metadata.inProgressSequences.get(node.nodeId()))
                                              .advance(metadata.epoch)
                                              .barrier()
                                              .withMessagingService(delivery);
            long before = TCMMetrics.instance.progressBarrierRetries.getCount();
            progressBarrier.await();
            Assert.assertTrue(TCMMetrics.instance.progressBarrierRetries.getCount() - before > 0);
            Assert.assertTrue(responded.size() == 1);
        }
    }

    // TODO: move to a common lib? Is this a good idea?
    public static Surjections.Surjection<Integer> ints(int minInclusive, int maxInclusive)
    {
        return l -> RngUtils.asInt(l, minInclusive, maxInclusive);
    }

    public static Surjections.Surjection<Long> longs()
    {
        return l -> PCGFastPure.next(l, l);
    }

    public static Surjections.Surjection<Boolean> bools()
    {
        return RngUtils::asBoolean;
    }

    public static <T1, T2, T3, T4, RES> Surjections.Surjection<RES> combine(Surjections.Surjection<T1> gen1, Surjections.Surjection<T2> gen2,
                                                                            Surjections.Surjection<T3> gen3, Surjections.Surjection<T4> gen4,
                                                                            IIsolatedExecutor.QuadFunction<T1, T2, T3, T4, RES> res)
    {
        return (long l) -> {
            return res.apply(gen1.inflate(PCGFastPure.next(l, 1)),
                             gen2.inflate(PCGFastPure.next(l, 2)),
                             gen3.inflate(PCGFastPure.next(l, 2)),
                             gen4.inflate(PCGFastPure.next(l, 2)));
        };
    }
}
