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

import java.util.ArrayList;
import java.util.Iterator;

import javax.jms.JMSException;
import javax.transaction.xa.XAException;

import org.activemq.message.ActiveMQMessage;
import org.activemq.message.ActiveMQXid;
import org.activemq.message.MessageAck;
import org.activemq.store.MessageStore;
import org.activemq.store.ProxyMessageStore;
import org.activemq.store.ProxyTopicMessageStore;
import org.activemq.store.TopicMessageStore;
import org.activemq.store.TransactionStore;

import EDU.oswego.cs.dl.util.concurrent.ConcurrentHashMap;

/**
 * Provides a TransactionStore implementation that can create transaction aware
 * MessageStore objects from non transaction aware MessageStore objects.
 * 
 * @version $Revision: 1.1.1.1 $
 */
public class VMTransactionStore implements TransactionStore {

    ConcurrentHashMap inflightTransactions = new ConcurrentHashMap();

    ConcurrentHashMap preparedTransactions = new ConcurrentHashMap();

    private boolean doingRecover;

    public static class Tx {
        private ArrayList messages = new ArrayList();

        private ArrayList acks = new ArrayList();

        public void add(AddMessageCommand msg) {
            messages.add(msg);
        }

        public void add(RemoveMessageCommand ack) {
            acks.add(ack);
        }

        public ActiveMQMessage[] getMessages() {
            ActiveMQMessage rc[] = new ActiveMQMessage[messages.size()];
            int count=0;
            for (Iterator iter = messages.iterator(); iter.hasNext();) {
                AddMessageCommand cmd = (AddMessageCommand) iter.next();
                rc[count++] = cmd.getMessage(); 
            }
            return rc;
        }

        public MessageAck[] getAcks() {
            MessageAck rc[] = new MessageAck[acks.size()];
            int count=0;
            for (Iterator iter = acks.iterator(); iter.hasNext();) {
                RemoveMessageCommand cmd = (RemoveMessageCommand) iter.next();
                rc[count++] = cmd.getMessageAck(); 
            }
            return rc;
        }

        /**
         * @throws JMSException
         */
        public void commit() throws XAException {
            try {
	            // Do all the message adds.            
	            for (Iterator iter = messages.iterator(); iter.hasNext();) {
	                AddMessageCommand cmd = (AddMessageCommand) iter.next();
	                cmd.run();
	            }
	            // And removes..
	            for (Iterator iter = acks.iterator(); iter.hasNext();) {
	                RemoveMessageCommand cmd = (RemoveMessageCommand) iter.next();
	                cmd.run();
	            }
            } catch ( JMSException e) {
                throw (XAException)new XAException(XAException.XAER_RMFAIL).initCause(e);
            }
        }
    }

    public interface AddMessageCommand {
        ActiveMQMessage getMessage();
        void run() throws JMSException;
    }

    public interface RemoveMessageCommand {
        MessageAck getMessageAck();
        void run() throws JMSException;
    }

    public MessageStore proxy(MessageStore messageStore) {
	    return new ProxyMessageStore(messageStore) {
	        public void addMessage(final ActiveMQMessage message) throws JMSException {
	            VMTransactionStore.this.addMessage(getDelegate(), message);
	        }
	
	        public void removeMessage(final MessageAck ack) throws JMSException {
	            VMTransactionStore.this.removeMessage(getDelegate(), ack);
	        }
	    };
    }

    public TopicMessageStore proxy(TopicMessageStore messageStore) {
	    return new ProxyTopicMessageStore(messageStore) {
	        public void addMessage(final ActiveMQMessage message) throws JMSException {
	            VMTransactionStore.this.addMessage(getDelegate(), message);
	        }
	        public void removeMessage(final MessageAck ack) throws JMSException {
	            VMTransactionStore.this.removeMessage(getDelegate(), ack);
	        }
	    };
    }

    /**
     * @see org.activemq.store.TransactionStore#prepare(org.activemq.service.Transaction)
     */
    public void prepare(Object txid) {
        Tx tx = (Tx) inflightTransactions.remove(txid);
        if (tx == null)
            return;
        preparedTransactions.put(txid, tx);
    }

    public Tx getTx(Object txid) {
        Tx tx = (Tx) inflightTransactions.get(txid);
        if (tx == null) {
            tx = new Tx();
            inflightTransactions.put(txid, tx);
        }
        return tx;
    }

    /**
     * @throws XAException
     * @see org.activemq.store.TransactionStore#commit(org.activemq.service.Transaction)
     */
    public void commit(Object txid, boolean wasPrepared) throws XAException {
        
        Tx tx;
        if( wasPrepared ) {
            tx = (Tx) preparedTransactions.remove(txid);
        } else {
            tx = (Tx) inflightTransactions.remove(txid);
        }
        
        if( tx == null )
            return;
        tx.commit();
        
    }

    /**
     * @see org.activemq.store.TransactionStore#rollback(org.activemq.service.Transaction)
     */
    public void rollback(Object txid) {
        preparedTransactions.remove(txid);
        inflightTransactions.remove(txid);
    }

    public void start() throws JMSException {
    }

    public void stop() throws JMSException {
    }

    synchronized public void recover(RecoveryListener listener) throws XAException {

        // All the inflight transactions get rolled back..
        inflightTransactions.clear();        
        this.doingRecover = true;
        try {
	        for (Iterator iter = preparedTransactions.keySet().iterator(); iter.hasNext();) {
	            Object txid = (Object) iter.next();
	            try {
	                Tx tx = (Tx) preparedTransactions.get(txid);
	                listener.recover((ActiveMQXid) txid, tx.getMessages(), tx.getAcks());
	            } catch (JMSException e) {
	                throw (XAException) new XAException("Recovery of a transaction failed:").initCause(e);
	            }
	        }
        } finally {
            this.doingRecover = false;
        }
    }

    /**
     * @param message
     * @throws JMSException
     */
    void addMessage(final MessageStore destination, final ActiveMQMessage message) throws JMSException {
        
        if( doingRecover )
            return;
        
        if (message.isPartOfTransaction()) {
            Tx tx = getTx(message.getTransactionId());
            tx.add(new AddMessageCommand() {
                public ActiveMQMessage getMessage() {
                    return message;
                }
                public void run() throws JMSException {
                    destination.addMessage(message);
                }
            });
        } else {
            destination.addMessage(message);
        }
    }
    
    /**
     * @param ack
     * @throws JMSException
     */
    private void removeMessage(final MessageStore destination,final MessageAck ack) throws JMSException {
        if( doingRecover )
            return;

        if (ack.isPartOfTransaction()) {
            Tx tx = getTx(ack.getTransactionId());
            tx.add(new RemoveMessageCommand() {
                public MessageAck getMessageAck() {
                    return ack;
                }
                public void run() throws JMSException {
                    destination.removeMessage(ack);
                }
            });
        } else {
            destination.removeMessage(ack);
        }
    }
    
}
