/**
 * 
 * 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.codehaus.activemq.service.impl;

import EDU.oswego.cs.dl.util.concurrent.ConcurrentHashMap;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.codehaus.activemq.broker.Broker;
import org.codehaus.activemq.broker.BrokerClient;
import org.codehaus.activemq.message.ActiveMQXid;
import org.codehaus.activemq.service.Transaction;
import org.codehaus.activemq.service.TransactionManager;
import org.codehaus.activemq.store.PreparedTransactionStore;
import org.codehaus.activemq.util.JMSExceptionHelper;

import javax.jms.JMSException;
import javax.transaction.xa.XAException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @version $Revision: 1.2 $
 */
public class TransactionManagerImpl implements TransactionManager {
    private static final Log log = LogFactory.getLog(TransactionManagerImpl.class);

    // the broker on which transactions operate
    private Broker broker;
    // The prepared XA transactions.
    private PreparedTransactionStore preparedTransactions;
    // Maps clients to the txids that they created.
    private Map activeClients = new ConcurrentHashMap();
    // Maps txids to ActiveMQTransactions
    private Map localTxs = new ConcurrentHashMap();
    // Maps txids to ActiveMQTransactions
    private Map xaTxs = new ConcurrentHashMap();
    // Keeps track of the current transaction
    private final ThreadLocal contextTx = new ThreadLocal();

    public TransactionManagerImpl(Broker broker, PreparedTransactionStore preparedTransactions) {
        this.preparedTransactions = preparedTransactions;
        this.broker = broker;
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#createLocalTransaction(org.codehaus.activemq.broker.BrokerClient, String)
     */
    public Transaction createLocalTransaction(final BrokerClient client, final String txid) throws JMSException {
        AbstractTransaction t = new LocalTransactionCommand(broker, localTxs, txid);
        localTxs.put(txid, t);
        return t;
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#createXATransaction(org.codehaus.activemq.broker.BrokerClient, org.codehaus.activemq.message.ActiveMQXid)
     */
    public Transaction createXATransaction(final BrokerClient client, final ActiveMQXid xid) throws XAException {
    	
    	// The xa transaction may allready be running.
    	Transaction tx = (Transaction) localTxs.get(xid);
    	if( tx == null ) {
        	if(log.isDebugEnabled())
        		log.debug("XA Transaction started: "+xid);
            tx = new XATransactionCommand(broker, xid, xaTxs, preparedTransactions);
            xaTxs.put(xid, tx);
    	}    		
        return tx;
        
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#getLocalTransaction(String)
     */
    public Transaction getLocalTransaction(String txid) throws JMSException {
        Transaction tx = (Transaction) localTxs.get(txid);
        if (tx == null) {
            throw new JMSException("Transaction '" + txid
                    + "' has not been started.");
        }
        return tx;
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#getXATransaction(org.codehaus.activemq.message.ActiveMQXid)
     */
    public Transaction getXATransaction(ActiveMQXid xid) throws XAException {
        Transaction tx = (Transaction) xaTxs.get(xid);
        if (tx == null) {
            XAException e = new XAException("Transaction '" + xid + "' has not been started.");
            e.errorCode = XAException.XAER_NOTA;
            throw e;
        }
        return tx;
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#getPreparedXATransactions()
     */
    public ActiveMQXid[] getPreparedXATransactions() throws XAException {
        return preparedTransactions.getXids();
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#setContexTransaction(org.codehaus.activemq.service.Transaction)
     */
    public void setContexTransaction(Transaction tx) {
        contextTx.set(tx);
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#getContexTransaction()
     */
    public Transaction getContexTransaction() {
        return (Transaction) contextTx.get();
    }

    /**
     * @see org.codehaus.activemq.service.TransactionManager#cleanUpClient(org.codehaus.activemq.broker.BrokerClient)
     */
    public void cleanUpClient(BrokerClient client) throws JMSException {
        // HRC: I don't think we need to keep track of the client's open transactions here...
        // It seems like BrokerClientImpl.close() allready cleans up open transactions.
        //
        List list = (List) activeClients.remove(client);
        if (list != null) {
            for (int i = 0; i < list.size(); i++) {
                try {
                    Object o = list.get(i);
                    if (o instanceof String) {
                        Transaction t = this.getLocalTransaction((String) o);
                        t.rollback();
                    }
                    else {
                        Transaction t = this.getXATransaction((ActiveMQXid) o);
                        t.rollback();
                    }
                }
                catch (Exception e) {
                    log.warn("ERROR Rolling back disconnected client's transactions: ", e);
                }
            }
            list.clear();
        }
    }

    public void loadTransaction(ActiveMQXid xid, Transaction transaction) throws XAException {
        // first lets associate any transient data structurs with the
        // transaction which has recently been loaded from disk
        if (transaction instanceof XATransactionCommand) {
            XATransactionCommand xaTransaction = (XATransactionCommand) transaction;
            xaTransaction.initialise(xaTxs, preparedTransactions);
        }
        transaction.setBroker(broker);

        xaTxs.put(xid, transaction);
    }

    public void start() throws JMSException {
        preparedTransactions.start();
        try {
            preparedTransactions.loadPreparedTransactions(this);
        }
        catch (XAException e) {
            throw JMSExceptionHelper.newJMSException("Failed to recover: " + e, e);
        }
    }

    public void stop() throws JMSException {
        preparedTransactions.stop();
    }



    // Implementation methods
    //-------------------------------------------------------------------------

    private void addActiveTransaction(BrokerClient client, Object transactionId) {
        List list = (List) activeClients.get(client);
        if (list == null) {
            list = new ArrayList();
            activeClients.put(client, list);
        }
        list.add(transactionId);
    }

    private void removeActiveTransaction(BrokerClient client, Object transactionId) {
        List list = (List) activeClients.get(client);
        if (list != null) {
            list.remove(transactionId);
        }
    }


}
