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

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.Vector;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import javax.jms.BytesMessage;
import javax.jms.Connection;
import javax.jms.ConnectionFactory;
import javax.jms.DeliveryMode;
import javax.jms.ExceptionListener;
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.Topic;
import javax.jms.TopicConnection;
import javax.jms.TopicConnectionFactory;
import javax.jms.TopicPublisher;
import javax.jms.TopicSession;
import javax.jms.TopicSubscriber;
import javax.management.ObjectName;
import junit.framework.Test;
import org.apache.activemq.ActiveMQConnection;
import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.CombinationTestSupport;
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.command.ActiveMQTopic;
import org.apache.activemq.store.kahadb.KahaDBPersistenceAdapter;
import org.apache.activemq.store.kahadb.KahaDBStore;
import org.apache.activemq.util.IOHelper;
import org.apache.activemq.util.Wait;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *  A Test case for AMQ-1479
 */
public class DurableConsumerTest extends CombinationTestSupport{
    private static final Logger LOG = LoggerFactory.getLogger(DurableConsumerTest.class);
    private static int COUNT = 1024;
    private static String CONSUMER_NAME = "DURABLE_TEST";
    protected BrokerService broker;
    
    protected String bindAddress = "tcp://localhost:61616";
    
    protected byte[] payload = new byte[1024 * 32];
    protected ConnectionFactory factory;
    protected Vector<Exception> exceptions = new Vector<Exception>();
    
    private static final String TOPIC_NAME = "failoverTopic";
    private static final String CONNECTION_URL = "failover:(tcp://localhost:61616,tcp://localhost:61617)";
    public boolean useDedicatedTaskRunner = false;
    
    private class SimpleTopicSubscriber implements MessageListener,ExceptionListener{
        
        private TopicConnection topicConnection = null;
        
        public SimpleTopicSubscriber(String connectionURL,String clientId,String topicName) {
            
            ActiveMQConnectionFactory topicConnectionFactory = null;
            TopicSession topicSession = null;
            Topic topic = null;
            TopicSubscriber topicSubscriber = null;
            
            topicConnectionFactory = new ActiveMQConnectionFactory(connectionURL);
            try {
                
                topic = new ActiveMQTopic(topicName);
                topicConnection = topicConnectionFactory.createTopicConnection();
                topicConnection.setClientID((clientId));
                topicConnection.start();
                
                topicSession = topicConnection.createTopicSession(false, Session.AUTO_ACKNOWLEDGE);
                topicSubscriber = topicSession.createDurableSubscriber(topic, (clientId));
                topicSubscriber.setMessageListener(this);
                
            } catch (JMSException e) {
                e.printStackTrace();
            }
        }
        
        public void onMessage(Message arg0){
        }
        
        public void closeConnection(){
            if (topicConnection != null) {
                try {
                    topicConnection.close();
                } catch (JMSException e) {
                }
            }
        }
        
        public void onException(JMSException exception){
            exceptions.add(exception);
        }
    }
    
    private class MessagePublisher implements Runnable{
        private final boolean shouldPublish = true;
        
        public void run(){
            TopicConnectionFactory topicConnectionFactory = null;
            TopicConnection topicConnection = null;
            TopicSession topicSession = null;
            Topic topic = null;
            TopicPublisher topicPublisher = null;
            Message message = null;
            
            topicConnectionFactory = new ActiveMQConnectionFactory(CONNECTION_URL);
            try {
                topic = new ActiveMQTopic(TOPIC_NAME);
                topicConnection = topicConnectionFactory.createTopicConnection();
                topicSession = topicConnection.createTopicSession(false, Session.AUTO_ACKNOWLEDGE);
                topicPublisher = topicSession.createPublisher(topic);
                message = topicSession.createMessage();
            } catch (Exception ex) {
                exceptions.add(ex);
            }
            while (shouldPublish) {
                try {
                    topicPublisher.publish(message, DeliveryMode.PERSISTENT, 1, 2 * 60 * 60 * 1000);
                } catch (JMSException ex) {
                    exceptions.add(ex);
                }
                try {
                    Thread.sleep(1);
                } catch (Exception ex) {
                }
            }
        }
    }
    
    private void configurePersistence(BrokerService broker) throws Exception{
        File dataDirFile = new File("target/" + getName());
        KahaDBPersistenceAdapter kahaDBAdapter = new KahaDBPersistenceAdapter();
        kahaDBAdapter.setDirectory(dataDirFile);
        broker.setPersistenceAdapter(kahaDBAdapter);
    }
    
    public void testFailover() throws Exception{
        
        configurePersistence(broker);
        broker.start();
        
        Thread publisherThread = new Thread(new MessagePublisher());
        publisherThread.start();
        final int numSubs = 100;
        final List<SimpleTopicSubscriber> list = new ArrayList<SimpleTopicSubscriber>(numSubs);
        for (int i = 0; i < numSubs; i++) {
            
            final int id = i;
            Thread thread = new Thread(new Runnable(){
                public void run(){
                    SimpleTopicSubscriber s =new SimpleTopicSubscriber(CONNECTION_URL, System.currentTimeMillis() + "-" + id, TOPIC_NAME);
                    list.add(s);
                }
            });
            thread.start();
            
        }

        Wait.waitFor(new Wait.Condition(){
            @Override
            public boolean isSatisified() throws Exception {
                return numSubs == list.size();
            }
        });

        broker.stop();
        broker = createBroker(false);
        configurePersistence(broker);
        broker.start();
        Thread.sleep(10000);
        for (SimpleTopicSubscriber s:list) {
            s.closeConnection();
        }
        assertTrue("no exceptions: " + exceptions, exceptions.isEmpty());
    }
    
    // makes heavy use of threads and can demonstrate https://issues.apache.org/activemq/browse/AMQ-2028
    // with use dedicatedTaskRunner=true and produce OOM
    public void initCombosForTestConcurrentDurableConsumer(){
        addCombinationValues("useDedicatedTaskRunner", new Object[] { Boolean.TRUE, Boolean.FALSE });
    }
    
    public void testConcurrentDurableConsumer() throws Exception{
        
        broker.start();
        broker.waitUntilStarted();
        
        factory = createConnectionFactory();
        final String topicName = getName();
        final int numMessages = 500;
        int numConsumers = 1;
        final CountDownLatch counsumerStarted = new CountDownLatch(numConsumers);
        final AtomicInteger receivedCount = new AtomicInteger();
        Runnable consumer = new Runnable(){
            public void run(){
                final String consumerName = Thread.currentThread().getName();
                int acked = 0;
                int received = 0;
                
                try {
                    while (acked < numMessages / 2) {
                        // take one message and close, ack on occasion
                        Connection consumerConnection = factory.createConnection();
                        ((ActiveMQConnection) consumerConnection).setWatchTopicAdvisories(false);
                        consumerConnection.setClientID(consumerName);
                        Session consumerSession = consumerConnection.createSession(false, Session.CLIENT_ACKNOWLEDGE);
                        Topic topic = consumerSession.createTopic(topicName);
                        consumerConnection.start();
                        
                        MessageConsumer consumer = consumerSession.createDurableSubscriber(topic, consumerName);
                        
                        counsumerStarted.countDown();
                        Message msg = null;
                        do {
                            msg = consumer.receive(5000);
                            if (msg != null) {
                                receivedCount.incrementAndGet();
                                if (received != 0 && received % 100 == 0) {
                                    LOG.info("Received msg: " + msg.getJMSMessageID());
                                }
                                if (++received % 2 == 0) {
                                    msg.acknowledge();
                                    acked++;
                                }
                            }
                        } while (msg == null);

                        consumerConnection.close();
                    }
                    assertTrue(received >= acked);
                } catch (Exception e) {
                    e.printStackTrace();
                    exceptions.add(e);
                }
            }
        };
        
        ExecutorService executor = Executors.newFixedThreadPool(numConsumers);
        
        for (int i = 0; i < numConsumers; i++) {
            executor.execute(consumer);
        }
        
        assertTrue(counsumerStarted.await(30, TimeUnit.SECONDS));
        
        Connection producerConnection = factory.createConnection();
        ((ActiveMQConnection) producerConnection).setWatchTopicAdvisories(false);
        Session producerSession = producerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        Topic topic = producerSession.createTopic(topicName);
        MessageProducer producer = producerSession.createProducer(topic);
        producerConnection.start();
        for (int i = 0; i < numMessages; i++) {
            BytesMessage msg = producerSession.createBytesMessage();
            msg.writeBytes(payload);
            producer.send(msg);
            if (i != 0 && i % 100 == 0) {
                LOG.info("Sent msg " + i);
            }
        }
        
        executor.shutdown();
        executor.awaitTermination(30, TimeUnit.SECONDS);
        
        Wait.waitFor(new Wait.Condition(){
            public boolean isSatisified() throws Exception{
                LOG.info("receivedCount: " + receivedCount.get());
                return receivedCount.get() == numMessages;
            }
        }, 360 * 1000);
        assertEquals("got required some messages", numMessages, receivedCount.get());
        assertTrue("no exceptions, but: " + exceptions, exceptions.isEmpty());
    }
    
    public void testConsumerRecover() throws Exception{
        doTestConsumer(true);
    }
    
    public void testConsumer() throws Exception{
        doTestConsumer(false);
    }

    public void testPrefetchViaBrokerConfig() throws Exception {

        Integer prefetchVal = new Integer(150);
        PolicyEntry policyEntry = new PolicyEntry();
        policyEntry.setDurableTopicPrefetch(prefetchVal.intValue());
        policyEntry.setPrioritizedMessages(true);
        PolicyMap policyMap = new PolicyMap();
        policyMap.setDefaultEntry(policyEntry);
        broker.setDestinationPolicy(policyMap);
        broker.start();

        factory = createConnectionFactory();
        Connection consumerConnection = factory.createConnection();
        consumerConnection.setClientID(CONSUMER_NAME);
        Session consumerSession = consumerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        Topic topic = consumerSession.createTopic(getClass().getName());
        MessageConsumer consumer = consumerSession.createDurableSubscriber(topic, CONSUMER_NAME);
        consumerConnection.start();

        ObjectName activeSubscriptionObjectName = broker.getAdminView().getDurableTopicSubscribers()[0];
        Object prefetchFromSubView = broker.getManagementContext().getAttribute(activeSubscriptionObjectName, "PrefetchSize");
        assertEquals(prefetchVal, prefetchFromSubView);
    }
    
    public void doTestConsumer(boolean forceRecover) throws Exception{
        
        if (forceRecover) {
            configurePersistence(broker);
        }
        broker.start();
        
        factory = createConnectionFactory();
        Connection consumerConnection = factory.createConnection();
        consumerConnection.setClientID(CONSUMER_NAME);
        Session consumerSession = consumerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        Topic topic = consumerSession.createTopic(getClass().getName());
        MessageConsumer consumer = consumerSession.createDurableSubscriber(topic, CONSUMER_NAME);
        consumerConnection.start();
        consumerConnection.close();
        broker.stop();
        broker = createBroker(false);
        if (forceRecover) {
            configurePersistence(broker);
        }
        broker.start();
        
        Connection producerConnection = factory.createConnection();
        
        Session producerSession = producerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        
        MessageProducer producer = producerSession.createProducer(topic);
        producerConnection.start();
        for (int i = 0; i < COUNT; i++) {
            BytesMessage msg = producerSession.createBytesMessage();
            msg.writeBytes(payload);
            producer.send(msg);
            if (i != 0 && i % 1000 == 0) {
                LOG.info("Sent msg " + i);
            }
        }
        producerConnection.close();
        broker.stop();
        broker = createBroker(false);
        if (forceRecover) {
            configurePersistence(broker);
        }
        broker.start();
        
        consumerConnection = factory.createConnection();
        consumerConnection.setClientID(CONSUMER_NAME);
        consumerConnection.start();
        consumerSession = consumerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        
        consumer = consumerSession.createDurableSubscriber(topic, CONSUMER_NAME);
        for (int i = 0; i < COUNT; i++) {
            Message msg = consumer.receive(10000);
            assertNotNull("Missing message: " + i, msg);
            if (i != 0 && i % 1000 == 0) {
                LOG.info("Received msg " + i);
            }
            
        }
        consumerConnection.close();
        
    }
    
    @Override
    protected void setUp() throws Exception{
        if (broker == null) {
            broker = createBroker(true);
        }
        
        super.setUp();
    }
    
    @Override
    protected void tearDown() throws Exception{
        super.tearDown();
        if (broker != null) {
            broker.stop();
            broker.waitUntilStopped();
            broker = null;
        }
    }
    
    protected Topic creatTopic(Session s,String destinationName) throws JMSException{
        return s.createTopic(destinationName);
    }
    
    /**
     * Factory method to create a new broker
     * 
     * @throws Exception
     */
    protected BrokerService createBroker(boolean deleteStore) throws Exception{
        BrokerService answer = new BrokerService();
        configureBroker(answer, deleteStore);
        return answer;
    }
    
    protected void configureBroker(BrokerService answer,boolean deleteStore) throws Exception{
        answer.setDeleteAllMessagesOnStartup(deleteStore);
        KahaDBStore kaha = new KahaDBStore();
        //kaha.setConcurrentStoreAndDispatchTopics(false);
        File directory = new File("target/activemq-data/kahadb");
        if (deleteStore) {
            IOHelper.deleteChildren(directory);
        }
        kaha.setDirectory(directory);
        //kaha.setMaxAsyncJobs(10);
        
        answer.setPersistenceAdapter(kaha);
        answer.addConnector(bindAddress);
        answer.setUseShutdownHook(false);
        answer.setAdvisorySupport(false);
        answer.setDedicatedTaskRunner(useDedicatedTaskRunner);
    }
    
    protected ActiveMQConnectionFactory createConnectionFactory() throws Exception{
        ActiveMQConnectionFactory factory = new ActiveMQConnectionFactory(bindAddress);
        factory.setUseDedicatedTaskRunner(useDedicatedTaskRunner);
        return factory;
    }
    
    public static Test suite(){
        return suite(DurableConsumerTest.class);
    }
    
    public static void main(String[] args){
        junit.textui.TestRunner.run(suite());
    }
}
