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

import java.net.URI;
import java.util.Collection;
import java.util.Iterator;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.jms.Connection;
import javax.jms.Destination;
import javax.jms.Message;
import javax.jms.MessageConsumer;
import javax.jms.Session;

import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.ActiveMQPrefetchPolicy;
import org.apache.activemq.JmsMultipleBrokersTestSupport;
import org.apache.activemq.broker.BrokerPlugin;
import org.apache.activemq.broker.BrokerPluginSupport;
import org.apache.activemq.broker.BrokerService;
import org.apache.activemq.broker.ProducerBrokerExchange;
import org.apache.activemq.broker.region.policy.PolicyEntry;
import org.apache.activemq.broker.region.policy.PolicyMap;
import org.apache.activemq.command.ActiveMQDestination;
import org.apache.activemq.network.ConditionalNetworkBridgeFilterFactory;
import org.apache.activemq.network.NetworkConnector;
import org.apache.activemq.util.Wait;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 */
public class TwoBrokerQueueClientsReconnectTest extends JmsMultipleBrokersTestSupport {
    protected static final int MESSAGE_COUNT = 100; // Best if a factor of 100
    protected static final int PREFETCH_COUNT = 1;
    protected static final int NETWORK_PREFETCH = 1;
    private static final Logger LOG = LoggerFactory.getLogger(TwoBrokerQueueClientsReconnectTest.class);


    protected int msgsClient1;
    protected int msgsClient2;
    protected String broker1;
    protected String broker2;

    public void testClientAReceivesOnly() throws Exception {
        broker1 = "BrokerA";
        broker2 = "BrokerB";

        doOneClientReceivesOnly();
    }

    public void testClientBReceivesOnly() throws Exception {
        broker1 = "BrokerB";
        broker2 = "BrokerA";

        doOneClientReceivesOnly();
    }

    public void doOneClientReceivesOnly() throws Exception {
        // allow immediate replay back to origin
        applyRateLimitNetworkFilter(0);

        // Bridge brokers
        bridgeBrokers(broker1, broker2);
        bridgeBrokers(broker2, broker1);

        // Run brokers
        startAllBrokers();

        // Create queue
        Destination dest = createDestination("TEST.FOO", false);

        // Create consumers
        MessageConsumer client1 = createConsumer(broker1, dest);
        MessageConsumer client2 = createConsumer(broker2, dest);

        // Give clients time to register with broker
        Thread.sleep(500);

        // Always send messages to broker A
        sendMessages("BrokerA", dest, MESSAGE_COUNT);

        // Close the second client, messages should be sent to the first client
        client2.close();

        // Let the first client receive all messages
        msgsClient1 += receiveAllMessages(client1);
        client1.close();

        // First client should have received 100 messages
        assertEquals("Client for " + broker1 + " should have receive all messages.", MESSAGE_COUNT, msgsClient1);
    }

    public void testClientAReceivesOnlyAfterReconnect() throws Exception {
        broker1 = "BrokerA";
        broker2 = "BrokerB";

        doOneClientReceivesOnlyAfterReconnect();
    }

    public void testClientBReceivesOnlyAfterReconnect() throws Exception {
        broker1 = "BrokerB";
        broker2 = "BrokerA";

        doOneClientReceivesOnlyAfterReconnect();
    }

    public void doOneClientReceivesOnlyAfterReconnect() throws Exception {
        // allow immediate replay back to origin
        applyRateLimitNetworkFilter(0);

        // Bridge brokers
        bridgeBrokers(broker1, broker2);
        bridgeBrokers(broker2, broker1);

        // Run brokers
        startAllBrokers();

        // Create queue
        Destination dest = createDestination("TEST.FOO", false);

        // Create first consumer
        MessageConsumer client1 = createConsumer(broker1, dest);
        MessageConsumer client2 = createConsumer(broker2, dest);

        // Give clients time to register with broker
        Thread.sleep(500);

        // Always send message to broker A
        sendMessages("BrokerA", dest, MESSAGE_COUNT);

        // Let the first client receive the first 20% of messages
        msgsClient1 += receiveExactMessages(client1, (int)(MESSAGE_COUNT * 0.20));

        // Disconnect the first client
        client1.close();

        // Create another client for the first broker
        client1 = createConsumer(broker1, dest);
        Thread.sleep(500);

        // Close the second client, messages should be sent to the first client
        client2.close();

        // Receive the rest of the messages
        msgsClient1 += receiveAllMessages(client1);
        client1.close();

        // The first client should have received 100 messages
        assertEquals("Client for " + broker1 + " should have received all messages.", MESSAGE_COUNT, msgsClient1);
    }

    public void testTwoClientsReceiveClientADisconnects() throws Exception {
        broker1 = "BrokerA";
        broker2 = "BrokerB";

        doTwoClientsReceiveOneClientDisconnects();
    }

    public void testTwoClientsReceiveClientBDisconnects() throws Exception {
        broker1 = "BrokerB";
        broker2 = "BrokerA";

        doTwoClientsReceiveOneClientDisconnects();
    }

    public void doTwoClientsReceiveOneClientDisconnects() throws Exception {
        // ensure all message do not flow across the network too quickly
        applyRateLimitNetworkFilter(0.8 * MESSAGE_COUNT);

        // Bridge brokers
        bridgeBrokers(broker1, broker2);
        bridgeBrokers(broker2, broker1);

        // Run brokers
        startAllBrokers();

        // Create queue
        Destination dest = createDestination("TEST.FOO", false);

        // Create first client
        MessageConsumer client1 = createConsumer(broker1, dest);
        MessageConsumer client2 = createConsumer(broker2, dest);

        // Give clients time to register with broker
        Thread.sleep(500);

        // Always send messages to broker A
        sendMessages("BrokerA", dest, MESSAGE_COUNT);

        LOG.info("Let each client receive 20% of the messages - 40% total");
        msgsClient1 += receiveExactMessages(client1, (int)(MESSAGE_COUNT * 0.20));
        msgsClient2 += receiveExactMessages(client2, (int)(MESSAGE_COUNT * 0.20));

        // Disconnect the first client
        client1.close();

        LOG.info("Let the second client receive the rest of the messages");
        msgsClient2 += receiveAllMessages(client2);
        client2.close();

        // First client should have received 20% of the messages
        assertEquals("Client for " + broker1 + " should have received 20% of the messages.", (int)(MESSAGE_COUNT * 0.20), msgsClient1);

        // Second client should have received 80% of the messages
        assertEquals("Client for " + broker2 + " should have received 80% of the messages.", (int)(MESSAGE_COUNT * 0.80), msgsClient2);
    }

    public void testTwoClientsReceiveClientAReconnects() throws Exception {
        broker1 = "BrokerA";
        broker2 = "BrokerB";

        doTwoClientsReceiveOneClientReconnects();
    }

    public void testTwoClientsReceiveClientBReconnects() throws Exception {
        broker1 = "BrokerB";
        broker2 = "BrokerA";

        doTwoClientsReceiveOneClientReconnects();
    }

    public void doTwoClientsReceiveOneClientReconnects() throws Exception {
        // ensure all message do not flow across the network too quickly
        applyRateLimitNetworkFilter(0.2 * MESSAGE_COUNT);

        // Bridge brokers
        bridgeBrokers(broker1, broker2);
        bridgeBrokers(broker2, broker1);

        // Run brokers
        startAllBrokers();

        // Create queue
        Destination dest = createDestination("TEST.FOO", false);

        // Create the first client
        MessageConsumer client1 = createConsumer(broker1, dest);
        MessageConsumer client2 = createConsumer(broker2, dest);

        // Give clients time to register with broker
        Thread.sleep(500);

        // Always send messages to broker A
        sendMessages("BrokerA", dest, MESSAGE_COUNT);

        // Let each client receive 20% of the messages - 40% total
        msgsClient1 += receiveExactMessages(client1, (int)(MESSAGE_COUNT * 0.20));
        msgsClient2 += receiveExactMessages(client2, (int)(MESSAGE_COUNT * 0.20));

        LOG.info("msgsClient1=" + msgsClient1);
        LOG.info("msgsClient2=" + msgsClient2);

        Thread.sleep(1000);
        LOG.info("Disconnect the first client");
        client1.close();

        LOG.info("Let the second client receive 20% more of the total messages");
        msgsClient2 += receiveExactMessages(client2, (int)(MESSAGE_COUNT * 0.20));

        LOG.info("msgsClient2=" + msgsClient2);

        // Create another client for broker 1
        client1 = createConsumer(broker1, dest);
        Thread.sleep(1000);

        // Let each client receive 20% of the messages - 40% total
        msgsClient1 += receiveExactMessages(client1, (int)(MESSAGE_COUNT * 0.20));
        client1.close();
        LOG.info("new consumer addition, msgsClient1=" + msgsClient1);

        Thread.sleep(2000);
        msgsClient2 += receiveExactMessages(client2, (int)(MESSAGE_COUNT * 0.20));
        client2.close();
        LOG.info("msgsClient2=" + msgsClient2);

        // First client should have received 40 messages
        assertEquals("Client for " + broker1 + " should have received 40% of the messages.", (int)(MESSAGE_COUNT * 0.40), msgsClient1);

        // Second client should have received 60 messages
        assertEquals("Client for " + broker2 + " should have received 60% of the messages.", (int)(MESSAGE_COUNT * 0.60), msgsClient2);
    }

    private void applyRateLimitNetworkFilter(double rateLimit) {
        ConditionalNetworkBridgeFilterFactory filterFactory = new ConditionalNetworkBridgeFilterFactory();
        filterFactory.setReplayWhenNoConsumers(true);
        filterFactory.setRateLimit((int) rateLimit);
        filterFactory.setRateDuration(1000);

        Collection<BrokerItem> brokerList = brokers.values();
        for (Iterator<BrokerItem> i = brokerList.iterator(); i.hasNext();) {
            BrokerService broker = i.next().broker;
            broker.getDestinationPolicy().getDefaultEntry().setNetworkBridgeFilterFactory(filterFactory);
        }
    }

    public void testTwoClientsReceiveTwoClientReconnects() throws Exception {
        // ensure all message do not flow across the network too quickly
        applyRateLimitNetworkFilter(0.5 * MESSAGE_COUNT);

        broker1 = "BrokerA";
        broker2 = "BrokerB";

        // Bridge brokers
        bridgeBrokers(broker1, broker2);
        bridgeBrokers(broker2, broker1);

        // Run brokers
        startAllBrokers();

        // Create queue
        Destination dest = createDestination("TEST.FOO", false);

        // Create the first client
        MessageConsumer client1 = createConsumer(broker1, dest);
        MessageConsumer client2 = createConsumer(broker2, dest);

        // Give clients time to register with broker
        Thread.sleep(500);

        // Always send messages to broker A
        sendMessages("BrokerA", dest, MESSAGE_COUNT);

        // Let each client receive 20% of the messages - 40% total
        msgsClient1 += receiveExactMessages(client1, (int)(MESSAGE_COUNT * 0.20));
        msgsClient2 += receiveExactMessages(client2, (int)(MESSAGE_COUNT * 0.20));

        LOG.info("Disconnect both clients");
        client1.close();
        client2.close();

        // Let each client receive 30% more of the total messages - 60% total
        LOG.info("Serially create another two clients for each broker and consume in turn");
        client1 = createConsumer(broker1, dest);
        msgsClient1 += receiveExactMessages(client1, (int)(MESSAGE_COUNT * 0.30));
        client1.close();

        // the close will allow replay or the replay of the remaining messages
        client2 = createConsumer(broker2, dest);
        msgsClient2 += receiveExactMessages(client2, (int)(MESSAGE_COUNT * 0.30));
        client2.close();

        // First client should have received 50% of the messages
        assertEquals("Client for " + broker1 + " should have received 50% of the messages.", (int)(MESSAGE_COUNT * 0.50), msgsClient1);

        // Second client should have received 50% of the messages
        assertEquals("Client for " + broker2 + " should have received 50% of the messages.", (int)(MESSAGE_COUNT * 0.50), msgsClient2);
    }

    @SuppressWarnings("unchecked")
    public void testDuplicateSend() throws Exception {
        broker1 = "BrokerA";
        broker2 = "BrokerB";

        // enable producer audit for the network connector, off by default b/c of interference with composite
        // dests and virtual topics
        brokers.get(broker2).broker.getTransportConnectors().get(0).setAuditNetworkProducers(true);
        bridgeBrokers(broker1, broker2);

        final AtomicBoolean first = new AtomicBoolean();
        final CountDownLatch gotMessageLatch = new CountDownLatch(1);

        BrokerService brokerService = brokers.get(broker2).broker;
        brokerService.setPersistent(true);
        brokerService.setDeleteAllMessagesOnStartup(true);
        brokerService.setPlugins(new BrokerPlugin[]{
                new BrokerPluginSupport() {
                    @Override
                    public void send(final ProducerBrokerExchange producerExchange,
                                     org.apache.activemq.command.Message messageSend)
                            throws Exception {
                        super.send(producerExchange, messageSend);
                        if (first.compareAndSet(false, true)) {
                            producerExchange.getConnectionContext().setDontSendReponse(true);
                            Executors.newSingleThreadExecutor().execute(new Runnable() {
                                @Override
                                public void run() {
                                    try {
                                        LOG.info("Waiting for recepit");
                                        assertTrue("message received on time", gotMessageLatch.await(60, TimeUnit.SECONDS));
                                        LOG.info("Stopping connection post send and receive and multiple producers");
                                        producerExchange.getConnectionContext().getConnection().stop();
                                    } catch (Exception e) {
                                        e.printStackTrace();
                                    }
                                }
                            });
                        }
                    }
                }
        });

        // Run brokers
        startAllBrokers();

        waitForBridgeFormation();

        // Create queue
        Destination dest = createDestination("TEST.FOO", false);

        MessageConsumer client2 = createConsumer(broker2, dest);

        sendMessages("BrokerA", dest, 1);

        assertEquals("Client got message", 1, receiveExactMessages(client2, 1));
        client2.close();
        gotMessageLatch.countDown();

        // message still pending on broker1
        assertEquals("messages message still there", 1, brokers.get(broker1).broker.getAdminView().getTotalMessageCount());

        client2 = createConsumer(broker2, dest);

        LOG.info("Let the second client receive the rest of the messages");
        assertEquals("no duplicate message", 0, receiveAllMessages(client2));
        assertEquals("no duplicate message", 0, receiveAllMessages(client2));

        assertEquals("no messages enqueued", 0, brokers.get(broker2).broker.getAdminView().getTotalMessageCount());
        assertTrue("no messages enqueued on origin", Wait.waitFor(new Wait.Condition() {
            @Override
            public boolean isSatisified() throws Exception {
                return 0 == brokers.get(broker1).broker.getAdminView().getTotalMessageCount();
            }
        }));
    }

    @SuppressWarnings("unchecked")
    public void testDuplicateSendWithNoAuditEnqueueCountStat() throws Exception {
        broker1 = "BrokerA";
        broker2 = "BrokerB";

        NetworkConnector networkConnector = bridgeBrokers(broker1, broker2);

        final AtomicBoolean first = new AtomicBoolean();
        final CountDownLatch gotMessageLatch = new CountDownLatch(1);

        BrokerService brokerService = brokers.get(broker2).broker;
        brokerService.setPersistent(true);
        brokerService.setDeleteAllMessagesOnStartup(true);
        brokerService.setPlugins(new BrokerPlugin[]{
                new BrokerPluginSupport() {
                    @Override
                    public void send(final ProducerBrokerExchange producerExchange,
                                     org.apache.activemq.command.Message messageSend)
                            throws Exception {
                        super.send(producerExchange, messageSend);
                        if (first.compareAndSet(false, true)) {
                            producerExchange.getConnectionContext().setDontSendReponse(true);
                            Executors.newSingleThreadExecutor().execute(new Runnable() {
                                @Override
                                public void run() {
                                    try {
                                        LOG.info("Waiting for recepit");
                                        assertTrue("message received on time", gotMessageLatch.await(60, TimeUnit.SECONDS));
                                        LOG.info("Stopping connection post send and receive and multiple producers");
                                        producerExchange.getConnectionContext().getConnection().stop();
                                    } catch (Exception e) {
                                        e.printStackTrace();
                                    }
                                }
                            });
                        }
                    }
                }
        });

        // Create queue
        ActiveMQDestination dest = createDestination("TEST.FOO", false);

        // statically include our destination
        networkConnector.addStaticallyIncludedDestination(dest);

        // Run brokers
        startAllBrokers();

        waitForBridgeFormation();

        sendMessages("BrokerA", dest, 1);

        // wait for broker2 to get the initial forward
        Wait.waitFor(new Wait.Condition(){
            @Override
            public boolean isSatisified() throws Exception {
                return brokers.get(broker2).broker.getAdminView().getTotalMessageCount() == 1;
            }
        });

        // message still pending on broker1
        assertEquals("messages message still there", 1, brokers.get(broker1).broker.getAdminView().getTotalMessageCount());

        // allow the bridge to be shutdown and restarted
        gotMessageLatch.countDown();


        // verify message is forwarded after restart
        assertTrue("no messages enqueued on origin", Wait.waitFor(new Wait.Condition() {
            @Override
            public boolean isSatisified() throws Exception {
                return 0 == brokers.get(broker1).broker.getAdminView().getTotalMessageCount();
            }
        }));

        assertEquals("one messages pending", 1, brokers.get(broker2).broker.getAdminView().getTotalMessageCount());
        assertEquals("one messages enqueued", 1, brokers.get(broker2).broker.getDestination(dest).getDestinationStatistics().getEnqueues().getCount());
    }

    protected int receiveExactMessages(MessageConsumer consumer, int msgCount) throws Exception {
        Message msg;
        int i;
        for (i = 0; i < msgCount; i++) {
            msg = consumer.receive(4000);
            if (msg == null) {
                LOG.error("Consumer failed to receive exactly " + msgCount + " messages. Actual messages received is: " + i);
                break;
            }
        }

        return i;
    }

    protected int receiveAllMessages(MessageConsumer consumer) throws Exception {
        int msgsReceived = 0;

        Message msg;
        do {
            msg = consumer.receive(1000);
            if (msg != null) {
                msgsReceived++;
            }
        } while (msg != null);

        return msgsReceived;
    }

    @Override
    protected MessageConsumer createConsumer(String brokerName, Destination dest) throws Exception {
        Connection conn = createConnection(brokerName);
        conn.start();
        Session sess = conn.createSession(false, Session.AUTO_ACKNOWLEDGE);
        return sess.createConsumer(dest);
    }

    @Override
    protected void configureBroker(BrokerService broker) {
        PolicyMap policyMap = new PolicyMap();
        PolicyEntry defaultEntry = new PolicyEntry();
        defaultEntry.setEnableAudit(false);
        policyMap.setDefaultEntry(defaultEntry);
        broker.setDestinationPolicy(policyMap);
    }

    @Override
    protected NetworkConnector bridgeBrokers(BrokerService localBroker, BrokerService remoteBroker, boolean dynamicOnly, int networkTTL, boolean conduit, boolean failover) throws Exception {
        NetworkConnector nc = super.bridgeBrokers(localBroker,remoteBroker, dynamicOnly, networkTTL, conduit, failover);
        nc.setPrefetchSize(NETWORK_PREFETCH);
        nc.setDecreaseNetworkConsumerPriority(true);
        return nc;
    }

    @Override
    public void setUp() throws Exception {
        super.setAutoFail(true);
        super.setUp();
        createBroker(new URI("broker:(tcp://localhost:61616)/BrokerA?persistent=false&useJmx=true"));
        createBroker(new URI("broker:(tcp://localhost:61617)/BrokerB?persistent=false&useJmx=true"));

        // Configure broker connection factory
        ActiveMQConnectionFactory factoryA;
        ActiveMQConnectionFactory factoryB;
        factoryA = (ActiveMQConnectionFactory)getConnectionFactory("BrokerA");
        factoryB = (ActiveMQConnectionFactory)getConnectionFactory("BrokerB");

        // Set prefetch policy
        ActiveMQPrefetchPolicy policy = new ActiveMQPrefetchPolicy();
        policy.setAll(PREFETCH_COUNT);

        factoryA.setPrefetchPolicy(policy);
        factoryB.setPrefetchPolicy(policy);

        msgsClient1 = 0;
        msgsClient2 = 0;
    }
}
