/**
 * 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.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import javax.jms.Connection;
import javax.jms.ConnectionFactory;
import javax.jms.DeliveryMode;
import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.MessageConsumer;
import javax.jms.MessageListener;
import javax.jms.MessageProducer;
import javax.jms.Session;
import javax.jms.TextMessage;
import javax.jms.TopicSubscriber;

import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.ActiveMQPrefetchPolicy;
import org.apache.activemq.TestSupport;
import org.apache.activemq.broker.BrokerService;
import org.apache.activemq.broker.region.policy.PolicyEntry;
import org.apache.activemq.broker.region.policy.PolicyMap;
import org.apache.activemq.broker.region.policy.StorePendingDurableSubscriberMessageStoragePolicy;
import org.apache.activemq.command.MessageId;
import org.apache.activemq.util.MessageIdList;
import org.apache.activemq.util.Wait;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
import org.junit.runners.Parameterized.Parameters;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@RunWith(value = Parameterized.class)
public class ConcurrentProducerDurableConsumerTest extends TestSupport {
    private static final Logger LOG = LoggerFactory.getLogger(ConcurrentProducerDurableConsumerTest.class);
    private final int consumerCount = 5;
    BrokerService broker;
    protected List<Connection> connections = Collections.synchronizedList(new ArrayList<Connection>());
    protected Map<MessageConsumer, TimedMessageListener> consumers = new HashMap<MessageConsumer, TimedMessageListener>();
    protected MessageIdList allMessagesList = new MessageIdList();
    private final int messageSize = 1024;

    private final TestSupport.PersistenceAdapterChoice persistenceAdapterChoice;

    @Parameters(name="{0}")
    public static Collection<TestSupport.PersistenceAdapterChoice[]> getTestParameters() {
        TestSupport.PersistenceAdapterChoice[] kahaDb = { TestSupport.PersistenceAdapterChoice.KahaDB };
        TestSupport.PersistenceAdapterChoice[] levelDb = { TestSupport.PersistenceAdapterChoice.LevelDB };
        TestSupport.PersistenceAdapterChoice[] mem = { TestSupport.PersistenceAdapterChoice.MEM };
        List<TestSupport.PersistenceAdapterChoice[]> choices = new ArrayList<TestSupport.PersistenceAdapterChoice[]>();
        choices.add(kahaDb);
        choices.add(levelDb);
        choices.add(mem);
        return choices;
    }

    public ConcurrentProducerDurableConsumerTest(TestSupport.PersistenceAdapterChoice choice) {
        this.persistenceAdapterChoice = choice;
    }

    @Test(timeout = 120000)
    public void testSendRateWithActivatingConsumers() throws Exception {
        final Destination destination = createDestination();
        final ConnectionFactory factory = createConnectionFactory();
        startInactiveConsumers(factory, destination);

        Connection connection = factory.createConnection();
        Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        MessageProducer producer = createMessageProducer(session, destination);

        // preload the durable consumers
        double[] inactiveConsumerStats = produceMessages(destination, 500, 10, session, producer, null);
        LOG.info("With inactive consumers: ave: " + inactiveConsumerStats[1] + ", max: " + inactiveConsumerStats[0] + ", multiplier: "
            + (inactiveConsumerStats[0] / inactiveConsumerStats[1]));

        // periodically start a durable sub that has a backlog
        final int consumersToActivate = 5;
        final Object addConsumerSignal = new Object();
        Executors.newCachedThreadPool(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                return new Thread(r, "ActivateConsumer" + this);
            }
        }).execute(new Runnable() {
            @Override
            public void run() {
                try {
                    MessageConsumer consumer = null;
                    for (int i = 0; i < consumersToActivate; i++) {
                        LOG.info("Waiting for add signal from producer...");
                        synchronized (addConsumerSignal) {
                            addConsumerSignal.wait(30 * 60 * 1000);
                        }
                        TimedMessageListener listener = new TimedMessageListener();
                        consumer = createDurableSubscriber(factory.createConnection(), destination, "consumer" + (i + 1));
                        LOG.info("Created consumer " + consumer);
                        consumer.setMessageListener(listener);
                        consumers.put(consumer, listener);
                    }
                } catch (Exception e) {
                    LOG.error("failed to start consumer", e);
                }
            }
        });

        double[] statsWithActive = produceMessages(destination, 500, 10, session, producer, addConsumerSignal);

        LOG.info(" with concurrent activate, ave: " + statsWithActive[1] + ", max: " + statsWithActive[0] + ", multiplier: "
            + (statsWithActive[0] / statsWithActive[1]));

        while (consumers.size() < consumersToActivate) {
            TimeUnit.SECONDS.sleep(2);
        }

        long timeToFirstAccumulator = 0;
        for (TimedMessageListener listener : consumers.values()) {
            long time = listener.getFirstReceipt();
            timeToFirstAccumulator += time;
            LOG.info("Time to first " + time);
        }
        LOG.info("Ave time to first message =" + timeToFirstAccumulator / consumers.size());

        for (TimedMessageListener listener : consumers.values()) {
            LOG.info("Ave batch receipt time: " + listener.waitForReceivedLimit(10000) + " max receipt: " + listener.maxReceiptTime);
        }

        // compare no active to active
        LOG.info("Ave send time with active: " + statsWithActive[1] + " as multiplier of ave with none active: " + inactiveConsumerStats[1] + ", multiplier="
            + (statsWithActive[1] / inactiveConsumerStats[1]));

        assertTrue("Ave send time with active: " + statsWithActive[1] + " within reasonable multpler of ave with none active: " + inactiveConsumerStats[1]
            + ", multiplier " + (statsWithActive[1] / inactiveConsumerStats[1]), statsWithActive[1] < 15 * inactiveConsumerStats[1]);
    }

    public void x_testSendWithInactiveAndActiveConsumers() throws Exception {
        Destination destination = createDestination();
        ConnectionFactory factory = createConnectionFactory();
        startInactiveConsumers(factory, destination);

        Connection connection = factory.createConnection();
        Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        MessageProducer producer = session.createProducer(destination);
        producer.setDeliveryMode(DeliveryMode.PERSISTENT);

        final int toSend = 100;
        final int numIterations = 5;

        double[] noConsumerStats = produceMessages(destination, toSend, numIterations, session, producer, null);

        startConsumers(factory, destination);
        LOG.info("Activated consumer");

        double[] withConsumerStats = produceMessages(destination, toSend, numIterations, session, producer, null);

        LOG.info("With consumer: " + withConsumerStats[1] + " , with noConsumer: " + noConsumerStats[1] + ", multiplier: "
            + (withConsumerStats[1] / noConsumerStats[1]));
        final int reasonableMultiplier = 15; // not so reasonable but improving
        assertTrue("max X times as slow with consumer: " + withConsumerStats[1] + ", with no Consumer: " + noConsumerStats[1] + ", multiplier: "
            + (withConsumerStats[1] / noConsumerStats[1]), withConsumerStats[1] < noConsumerStats[1] * reasonableMultiplier);

        final int toReceive = toSend * numIterations * consumerCount * 2;
        Wait.waitFor(new Wait.Condition() {
            @Override
            public boolean isSatisified() throws Exception {
                LOG.info("count: " + allMessagesList.getMessageCount());
                return toReceive == allMessagesList.getMessageCount();
            }
        }, 60 * 1000);

        assertEquals("got all messages", toReceive, allMessagesList.getMessageCount());
    }

    private MessageProducer createMessageProducer(Session session, Destination destination) throws JMSException {
        MessageProducer producer = session.createProducer(destination);
        producer.setDeliveryMode(DeliveryMode.PERSISTENT);
        return producer;
    }

    private void startInactiveConsumers(ConnectionFactory factory, Destination destination) throws Exception {
        // create off line consumers
        startConsumers(factory, destination);
        for (Connection connection : connections) {
            connection.close();
        }
        connections.clear();
        consumers.clear();
    }

    protected void startConsumers(ConnectionFactory factory, Destination dest) throws Exception {
        MessageConsumer consumer;
        for (int i = 0; i < consumerCount; i++) {
            TimedMessageListener list = new TimedMessageListener();
            consumer = createDurableSubscriber(factory.createConnection(), dest, "consumer" + (i + 1));
            consumer.setMessageListener(list);
            consumers.put(consumer, list);
        }
    }

    protected TopicSubscriber createDurableSubscriber(Connection conn, Destination dest, String name) throws Exception {
        conn.setClientID(name);
        connections.add(conn);
        conn.start();

        Session sess = conn.createSession(false, Session.AUTO_ACKNOWLEDGE);
        final TopicSubscriber consumer = sess.createDurableSubscriber((javax.jms.Topic) dest, name);

        return consumer;
    }

    /**
     * @return max and ave send time
     * @throws Exception
     */
    private double[] produceMessages(Destination destination, final int toSend, final int numIterations, Session session, MessageProducer producer,
        Object addConsumerSignal) throws Exception {
        long start;
        long count = 0;
        double batchMax = 0, max = 0, sum = 0;
        for (int i = 0; i < numIterations; i++) {
            start = System.currentTimeMillis();
            for (int j = 0; j < toSend; j++) {
                long singleSendstart = System.currentTimeMillis();
                TextMessage msg = createTextMessage(session, "" + j);
                // rotate
                int priority = ((int) count % 10);
                producer.send(msg, DeliveryMode.PERSISTENT, priority, 0);
                max = Math.max(max, (System.currentTimeMillis() - singleSendstart));
                if (++count % 500 == 0) {
                    if (addConsumerSignal != null) {
                        synchronized (addConsumerSignal) {
                            addConsumerSignal.notifyAll();
                            LOG.info("Signalled add consumer");
                        }
                    }
                }
                ;
                if (count % 5000 == 0) {
                    LOG.info("Sent " + count + ", singleSendMax:" + max);
                }

            }
            long duration = System.currentTimeMillis() - start;
            batchMax = Math.max(batchMax, duration);
            sum += duration;
            LOG.info("Iteration " + i + ", sent " + toSend + ", time: " + duration + ", batchMax:" + batchMax + ", singleSendMax:" + max);
        }

        LOG.info("Sent: " + toSend * numIterations + ", batchMax: " + batchMax + " singleSendMax: " + max);
        return new double[] { batchMax, sum / numIterations };
    }

    protected TextMessage createTextMessage(Session session, String initText) throws Exception {
        TextMessage msg = session.createTextMessage();

        // Pad message text
        if (initText.length() < messageSize) {
            char[] data = new char[messageSize - initText.length()];
            Arrays.fill(data, '*');
            String str = new String(data);
            msg.setText(initText + str);

            // Do not pad message text
        } else {
            msg.setText(initText);
        }

        return msg;
    }

    @Override
    @Before
    public void setUp() throws Exception {
        topic = true;
        super.setUp();
        broker = createBroker();
        broker.start();
    }

    @Override
    @After
    public void tearDown() throws Exception {
        for (Iterator<Connection> iter = connections.iterator(); iter.hasNext();) {
            Connection conn = iter.next();
            try {
                conn.close();
            } catch (Throwable e) {
            }
        }
        broker.stop();
        allMessagesList.flushMessages();
        consumers.clear();
        super.tearDown();
    }

    protected BrokerService createBroker() throws Exception {
        BrokerService brokerService = new BrokerService();
        brokerService.setEnableStatistics(false);
        brokerService.addConnector("tcp://0.0.0.0:0");
        brokerService.setDeleteAllMessagesOnStartup(true);

        PolicyEntry policy = new PolicyEntry();
        policy.setPrioritizedMessages(true);
        policy.setMaxPageSize(500);

        StorePendingDurableSubscriberMessageStoragePolicy durableSubPending = new StorePendingDurableSubscriberMessageStoragePolicy();
        durableSubPending.setImmediatePriorityDispatch(true);
        durableSubPending.setUseCache(true);
        policy.setPendingDurableSubscriberPolicy(durableSubPending);

        PolicyMap policyMap = new PolicyMap();
        policyMap.setDefaultEntry(policy);
        brokerService.setDestinationPolicy(policyMap);

        setPersistenceAdapter(brokerService, persistenceAdapterChoice);

        return brokerService;
    }

    @Override
    protected ActiveMQConnectionFactory createConnectionFactory() throws Exception {
        ActiveMQConnectionFactory factory = new ActiveMQConnectionFactory(broker.getTransportConnectors().get(0).getPublishableConnectString());
        ActiveMQPrefetchPolicy prefetchPolicy = new ActiveMQPrefetchPolicy();
        prefetchPolicy.setAll(1);
        factory.setPrefetchPolicy(prefetchPolicy);

        factory.setDispatchAsync(true);
        return factory;
    }

    class TimedMessageListener implements MessageListener {
        final int batchSize = 1000;
        CountDownLatch firstReceiptLatch = new CountDownLatch(1);
        long mark = System.currentTimeMillis();
        long firstReceipt = 0l;
        long receiptAccumulator = 0;
        long batchReceiptAccumulator = 0;
        long maxReceiptTime = 0;
        AtomicLong count = new AtomicLong(0);
        Map<Integer, MessageIdList> messageLists = new ConcurrentHashMap<Integer, MessageIdList>(new HashMap<Integer, MessageIdList>());

        @Override
        public void onMessage(Message message) {
            final long current = System.currentTimeMillis();
            final long duration = current - mark;
            receiptAccumulator += duration;
            int priority = 0;
            try {
                priority = message.getJMSPriority();
            } catch (JMSException ignored) {
            }
            if (!messageLists.containsKey(priority)) {
                MessageIdList perPriorityList = new MessageIdList();
                perPriorityList.setParent(allMessagesList);
                messageLists.put(priority, perPriorityList);
            }
            messageLists.get(priority).onMessage(message);
            if (count.incrementAndGet() == 1) {
                firstReceipt = duration;
                firstReceiptLatch.countDown();
                LOG.info("First receipt in " + firstReceipt + "ms");
            } else if (count.get() % batchSize == 0) {
                LOG.info("Consumed " + count.get() + " in " + batchReceiptAccumulator + "ms" + ", priority:" + priority);
                batchReceiptAccumulator = 0;
            }
            maxReceiptTime = Math.max(maxReceiptTime, duration);
            receiptAccumulator += duration;
            batchReceiptAccumulator += duration;
            mark = current;
        }

        long getMessageCount() {
            return count.get();
        }

        long getFirstReceipt() throws Exception {
            firstReceiptLatch.await(30, TimeUnit.SECONDS);
            return firstReceipt;
        }

        public long waitForReceivedLimit(long limit) throws Exception {
            final long expiry = System.currentTimeMillis() + 30 * 60 * 1000;
            while (count.get() < limit) {
                if (System.currentTimeMillis() > expiry) {
                    throw new RuntimeException("Expired waiting for X messages, " + limit);
                }
                TimeUnit.SECONDS.sleep(2);
                String missing = findFirstMissingMessage();
                if (missing != null) {
                    LOG.info("first missing = " + missing);
                    throw new RuntimeException("We have a missing message. " + missing);
                }

            }
            return receiptAccumulator / (limit / batchSize);
        }

        private String findFirstMissingMessage() {
            MessageId current = new MessageId();
            for (MessageIdList priorityList : messageLists.values()) {
                MessageId previous = null;
                for (String id : priorityList.getMessageIds()) {
                    current.setValue(id);
                    if (previous == null) {
                        previous = current.copy();
                    } else {
                        if (current.getProducerSequenceId() - 1 != previous.getProducerSequenceId()
                            && current.getProducerSequenceId() - 10 != previous.getProducerSequenceId()) {
                            return "Missing next after: " + previous + ", got: " + current;
                        } else {
                            previous = current.copy();
                        }
                    }
                }
            }
            return null;
        }
    }
}
