/** 
 * 
 * Copyright 2004 Protique Ltd
 * 
 * Licensed 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.activemq;
import javax.jms.BytesMessage;
import javax.jms.Connection;
import javax.jms.ConnectionFactory;
import javax.jms.DeliveryMode;
import javax.jms.Destination;
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 java.util.Iterator;
import java.util.List;

import junit.framework.TestCase;
import org.activemq.broker.BrokerContainer;
import org.activemq.broker.impl.BrokerContainerImpl;
import org.activemq.broker.impl.DefaultBroker;
import org.activemq.io.util.MemoryBoundedQueue;
import org.activemq.message.ActiveMQBytesMessage;
import org.activemq.message.ActiveMQQueue;
import org.activemq.message.ActiveMQTopic;
import org.activemq.util.IdGenerator;
import EDU.oswego.cs.dl.util.concurrent.SynchronizedInt;

/**
 * @version $Revision: 1.1.1.1 $
 */
public abstract class LargeMessageTestSupport extends TestCase implements MessageListener {
    protected static final int LARGE_MESSAGE_SIZE = 128 * 1024;
    protected static final int MESSAGE_COUNT = 100;
    protected BrokerContainer broker;
    protected Connection producerConnection;
    protected Connection consumerConnection;
    protected MessageConsumer consumer;
    protected MessageProducer producer;
    protected Session producerSession;
    protected Session consumerSession;
    protected byte[] largeMessageData;
    protected Destination destination;
    protected boolean isTopic = true;
    protected boolean isDurable = true;
    protected int deliveryMode = DeliveryMode.PERSISTENT;
    protected IdGenerator idGen = new IdGenerator();
    protected boolean validMessageConsumption = true;
    protected SynchronizedInt messageCount = new SynchronizedInt(0);
    protected String URL = "tcp://localhost:61717";
    protected DefaultBroker defaultBroker;

    protected int prefetchValue = 10000000;

    protected Destination createDestination() {
        String subject = getClass().getName();
        if (isTopic) {
            return new ActiveMQTopic(subject);
        }
        else {
            return new ActiveMQQueue(subject);
        }
    }

    protected MessageConsumer createConsumer() throws Exception {
        if (isTopic && isDurable) {
            return consumerSession.createDurableSubscriber((Topic) destination, idGen.generateId());
        }
        else {
            return consumerSession.createConsumer(destination);
        }
    }

    protected abstract ActiveMQConnectionFactory createConnectionFactory();

    protected void setUp() throws Exception {
        messageCount.set(0);
        broker = new BrokerContainerImpl("TestLargeMessages");
        broker.addConnector(URL);
        broker.start();
        destination = createDestination();
        largeMessageData = new byte[LARGE_MESSAGE_SIZE];
        for (int i = 0;i < LARGE_MESSAGE_SIZE;i++) {
            if (i % 2 == 0) {
                largeMessageData[i] = 'a';
            }
            else {
                largeMessageData[i] = 'z';
            }
        }
        defaultBroker = (DefaultBroker) broker.getBroker();
        Thread.sleep(1000);//allow the broker to start
        ActiveMQConnectionFactory fac = createConnectionFactory();
        fac.setQuickClose(true);
        producerConnection = fac.createConnection();
        setPrefetchPolicy((ActiveMQConnection) producerConnection);
        producerConnection.start();

        consumerConnection = fac.createConnection();
        setPrefetchPolicy((ActiveMQConnection) consumerConnection);
        consumerConnection.setClientID(idGen.generateId());
        consumerConnection.start();
        producerSession = producerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        producer = producerSession.createProducer(createDestination());
        producer.setDeliveryMode(deliveryMode);
        consumerSession = consumerConnection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        consumer = createConsumer();
        consumer.setMessageListener(this);
    }

    protected void setPrefetchPolicy(ActiveMQConnection activeMQConnection) {
        activeMQConnection.getPrefetchPolicy().setTopicPrefetch(prefetchValue);
        activeMQConnection.getPrefetchPolicy().setQueuePrefetch(prefetchValue);
        activeMQConnection.getPrefetchPolicy().setDurableTopicPrefetch(prefetchValue);
        activeMQConnection.getPrefetchPolicy().setQueueBrowserPrefetch(prefetchValue);
    }

    protected void tearDown() throws Exception {
        Thread.sleep(1000);
        producerConnection.close();
        consumerConnection.close();
        broker.stop();
        largeMessageData = null;
    }

    protected boolean isSame(BytesMessage msg1) throws Exception {
        boolean result = false;
        for (int i = 0;i < LARGE_MESSAGE_SIZE;i++) {
            result = msg1.readByte() == largeMessageData[i];
            if (!result)
                break;
        }
        return result;
    }

    public void onMessage(Message msg) {
        try {
            BytesMessage ba = (BytesMessage) msg;
            validMessageConsumption &= isSame(ba);
            if (messageCount.increment() >= MESSAGE_COUNT) {
                synchronized (messageCount) {
                    messageCount.notify();
                }
            }
            if (messageCount.get() % 50 == 0) {
                System.out.println("count = " + messageCount);
                System.out.println(defaultBroker.getQueueManager().dumpContents());
            }
        }
        catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void testLargeMessages() throws Exception {
        for (int i = 0;i < MESSAGE_COUNT;i++) {
            BytesMessage msg = producerSession.createBytesMessage();
            msg.writeBytes(largeMessageData);
            producer.send(msg);
        }
        synchronized (messageCount) {
            if (messageCount.get() < MESSAGE_COUNT) {
                messageCount.wait(60000);
            }
        }
        System.out.println("Finished count = " + messageCount );
        System.out.println(defaultBroker.getQueueManager().dumpContents());
        assertTrue("Not enough messages - expected " + MESSAGE_COUNT + " but got " + messageCount,
                messageCount.get() == MESSAGE_COUNT);
        assertTrue("received messages are not valid", validMessageConsumption);
        Thread.sleep(1000);
        System.out.println("FINAL count = " + messageCount);
        System.out.println(defaultBroker.getQueueManager().dumpContents());
        List list = defaultBroker.getQueueManager().getMemoryBoundedQueues();
        for (Iterator i = list.iterator(); i.hasNext();){
            MemoryBoundedQueue q = (MemoryBoundedQueue)i.next();
            if (q.getLocalMemoryUsedByThisQueue() > 0){
                System.err.println("Queue with enqueued messages at end of test: " + q.getName() + "= " + q.getLocalMemoryUsedByThisQueue());
            }
        }
    }
}