/**
 *
 * Copyright 2004 Protique Ltd
 * Copyright 2004 Hiram Chirino
 *
 * 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 java.util.ArrayList;

import javax.jms.JMSException;
import javax.jms.TransactionInProgressException;
import javax.jms.TransactionRolledBackException;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.activemq.message.ActiveMQXid;
import org.activemq.message.IntResponseReceipt;
import org.activemq.message.ResponseReceipt;
import org.activemq.message.TransactionInfo;
import org.activemq.message.XATransactionInfo;
import org.activemq.util.IdGenerator;

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

/**
 * A TransactionContext provides the means to control a JMS transaction.  It provides
 * a local transaction interface and also an XAResource interface.
 * 
 * <p/>
 * An application server controls the transactional assignment of an XASession
 * by obtaining its XAResource. It uses the XAResource to assign the session
 * to a transaction, prepare and commit work on the transaction, and so on.
 * <p/>
 * An XAResource provides some fairly sophisticated facilities for
 * interleaving work on multiple transactions, recovering a list of
 * transactions in progress, and so on. A JTA aware JMS provider must fully
 * implement this functionality. This could be done by using the services of a
 * database that supports XA, or a JMS provider may choose to implement this
 * functionality from scratch.
 * <p/>
 *
 * @version $Revision: 1.1.1.1 $
 * @see javax.jms.Session
 * @see javax.jms.QueueSession
 * @see javax.jms.TopicSession
 * @see javax.jms.XASession
 */
public class TransactionContext implements XAResource {
    
    private static final Log log = LogFactory.getLog(TransactionContext.class);
    
    private final ActiveMQConnection connection;
    private final ArrayList sessions = new ArrayList(2);
    private final SynchronizedBoolean startTransaction = new SynchronizedBoolean(false);
    private final IdGenerator localTransactionIdGenerator = new IdGenerator();

    // To track XA transactions.
    private Xid associatedXid;
    private ActiveMQXid activeXid;

    // To track local transactions.
    private String localTransactionId;

    private LocalTransactionEventListener localTransactionEventListener;
    
    public TransactionContext(ActiveMQConnection connection) {
        this.connection = connection;        
    }
    
    public boolean isInXATransaction() {
        return associatedXid!=null;
    }
    
    public boolean isInLocalTransaction() {
        return localTransactionId!=null;
    }
    
    /**
     * @return Returns the localTransactionEventListener.
     */
    public LocalTransactionEventListener getLocalTransactionEventListener() {
        return localTransactionEventListener;
    }

    /**
     * Used by the resource adapter to listen to transaction events.
     * 
     * @param localTransactionEventListener The localTransactionEventListener to set.
     */
    public void setLocalTransactionEventListener(LocalTransactionEventListener localTransactionEventListener) {
        this.localTransactionEventListener = localTransactionEventListener;
    }

    /////////////////////////////////////////////////////////////
    //
    // Methods that interface with the session
    //
    /////////////////////////////////////////////////////////////
    public void addSession(ActiveMQSession session) {
        sessions.add(session);
    }
    public void removeSession(ActiveMQSession session) {
        sessions.remove(session);
    }
    
    private void postRollback() {
        int size = sessions.size();        
        for(int i=0; i < size; i++ ){
            ((ActiveMQSession)sessions.get(i)).redeliverUnacknowledgedMessages(true);
        }
    }

    private void postCommit() {
        int size = sessions.size();        
        for(int i=0; i < size; i++ ){
            ((ActiveMQSession)sessions.get(i)).clearDeliveredMessages();
        }
    }

    public Object getTransactionId() {
        if( localTransactionId!=null )
            return localTransactionId;
        return activeXid;
    }
        
    /////////////////////////////////////////////////////////////
    //
    // Local transaction interface.
    //
    /////////////////////////////////////////////////////////////        
    
    /**
     * Start a local transaction.
     */
    public void begin() throws JMSException {        
        if( associatedXid!=null ) 
            throw new TransactionInProgressException("Cannot start local transction.  XA transaction is allready in progress.");        
        
        if( localTransactionId==null ) {         
            this.localTransactionId = localTransactionIdGenerator.generateId();
            TransactionInfo info = new TransactionInfo();
            info.setTransactionId((String)localTransactionId);
            info.setType(TransactionInfo.START);
            this.connection.asyncSendPacket(info);
            
            // Notify the listener that the tx was started.
            if (localTransactionEventListener != null) {
                localTransactionEventListener.beginEvent();
            }
        }
    }
    
    /**
     * Rolls back any messages done in this transaction and releases any locks currently held.
     * 
     * @throws JMSException if the JMS provider fails to roll back the transaction due to some internal error.
     * @throws javax.jms.IllegalStateException if the method is not called by a transacted session.
     */
    public void rollback() throws JMSException {
        if( associatedXid!=null ) 
            throw new TransactionInProgressException("Cannot rollback() if an XA transaction is allready in progress ");
        
        if( localTransactionId!=null ) {            
            TransactionInfo info = new TransactionInfo();
            info.setTransactionId((String)localTransactionId);
            info.setType(TransactionInfo.ROLLBACK);
            //before we send, update the current transaction id
            this.localTransactionId = null;
            this.connection.asyncSendPacket(info);
            // Notify the listener that the tx was rolled back
            if (localTransactionEventListener != null) {
                localTransactionEventListener.rollbackEvent();
            }
        }
        postRollback();
    }
        
    /**
     * Commits all messages done in this transaction and releases any locks currently held.
     * 
     * @throws JMSException if the JMS provider fails to commit the transaction due to some internal error.
     * @throws TransactionRolledBackException if the transaction is rolled back due to some internal error during
     * commit.
     * @throws javax.jms.IllegalStateException if the method is not called by a transacted session.
     */
    public void commit() throws JMSException {
        if( associatedXid!=null ) 
            throw new TransactionInProgressException("Cannot commit() if an XA transaction is allready in progress ");

        // Only send commit if the transaction was started.
        if (localTransactionId!=null) {
            TransactionInfo info = new TransactionInfo();
            info.setTransactionId((String)localTransactionId);
            info.setType(TransactionInfo.COMMIT);
            //before we send, update the current transaction id
            this.localTransactionId = null;
            // Notify the listener that the tx was commited back
            this.connection.syncSendPacket(info);
            if (localTransactionEventListener != null) {
                localTransactionEventListener.commitEvent();
            }
        }
        postCommit();
    }
    
    /////////////////////////////////////////////////////////////
    //
    // XAResource Implementation
    //
    /////////////////////////////////////////////////////////////
    /**
     * Associates a transaction with the resource.
     */
    public void start(Xid xid, int flags) throws XAException {
        if( localTransactionId!=null ) 
            throw new XAException(XAException.XAER_PROTO);
        
        // Are we allready associated?
        if (associatedXid != null) {
            throw new XAException(XAException.XAER_PROTO);
        }

        if ((flags & TMJOIN) == TMJOIN) {
            // TODO: verify that the server has seen the xid
        }
        if ((flags & TMJOIN) == TMRESUME) {
            // TODO: verify that the xid was suspended.
        }

        // associate
        setXid(xid);

        XATransactionInfo info = new XATransactionInfo();
        info.setXid(activeXid);
        info.setType(XATransactionInfo.START);

        try {
            // TODO: we may want to wait for reply..
            // server could fail this request
            this.connection.syncSendPacket(info);

        } catch (JMSException e) {
            throw toXAException(e);
        }
    }

    public void end(Xid xid, int flags) throws XAException {
        if( localTransactionId!=null ) 
            throw new XAException(XAException.XAER_PROTO);

        if ((flags & TMSUSPEND) == TMSUSPEND) {
            // You can only suspend the associated xid.
            if (associatedXid == null || !ActiveMQXid.equals(associatedXid,xid)) {
                throw new XAException(XAException.XAER_PROTO);
            }

            //TODO: we may want to put the xid in a suspended list.
            setXid(null);
        } else if ((flags & TMFAIL) == TMFAIL) {
            //TODO: We need to rollback the transaction??
            setXid(null);
        } else if ((flags & TMSUCCESS) == TMSUCCESS) {
            //set to null if this is the current xid.
            //otherwise this could be an asynchronous success call
            if (ActiveMQXid.equals(associatedXid,xid)) {
                setXid(null);
            }
        } else {
            throw new XAException(XAException.XAER_INVAL);
        }

    }

    public int prepare(Xid xid) throws XAException {

        // We allow interleaving multiple transactions, so
        // we don't limit prepare to the associated xid.
        ActiveMQXid x;
        //THIS SHOULD NEVER HAPPEN because end(xid, TMSUCCESS) should have been called first
        if (ActiveMQXid.equals(associatedXid,xid)) {
            throw new XAException(XAException.XAER_PROTO);
        } else {
            //TODO cache the known xids so we don't keep recreating this one??
            x = new ActiveMQXid(xid);
        }

        XATransactionInfo info = new XATransactionInfo();
        info.setXid(x);
        info.setType(XATransactionInfo.PRE_COMMIT);

        try {
            // Find out if the server wants to commit or rollback.
            IntResponseReceipt receipt = (IntResponseReceipt) this.connection.syncSendRequest(info);
            return receipt.getResult();
        } catch (JMSException e) {
            throw toXAException(e);
        }
    }

    public void rollback(Xid xid) throws XAException {

        // We allow interleaving multiple transactions, so
        // we don't limit rollback to the associated xid.
        ActiveMQXid x;
        if (ActiveMQXid.equals(associatedXid,xid)) {
            //I think this can happen even without an end(xid) call.  Need to check spec.
            x = activeXid;
        } else {
            x = new ActiveMQXid(xid);
        }

        XATransactionInfo info = new XATransactionInfo();
        info.setXid(x);
        info.setType(XATransactionInfo.ROLLBACK);

        try {
            // Let the server know that the tx is rollback.
            this.connection.syncSendPacket(info);
        } catch (JMSException e) {
            throw toXAException(e);
        }
        
        postRollback();
    }

    // XAResource interface
    public void commit(Xid xid, boolean onePhase) throws XAException {

        // We allow interleaving multiple transactions, so
        // we don't limit commit to the associated xid.
        ActiveMQXid x;
        if (ActiveMQXid.equals(associatedXid,xid)) {
            //should never happen, end(xid,TMSUCCESS) must have been previously called
            throw new XAException(XAException.XAER_PROTO);
        } else {
            x = new ActiveMQXid(xid);
        }

        XATransactionInfo info = new XATransactionInfo();
        info.setXid(x);
        info.setType(onePhase ? XATransactionInfo.COMMIT_ONE_PHASE : XATransactionInfo.COMMIT);

        try {
            // Notify the server that the tx was commited back
            this.connection.syncSendPacket(info);
        } catch (JMSException e) {
            throw toXAException(e);
        }

        postCommit();
    }


    public void forget(Xid xid) throws XAException {

        // We allow interleaving multiple transactions, so
        // we don't limit forget to the associated xid.
        ActiveMQXid x;
        if (ActiveMQXid.equals(associatedXid,xid)) {
            //TODO determine if this can happen... I think not.
            x = activeXid;
        } else {
            x = new ActiveMQXid(xid);
        }

        XATransactionInfo info = new XATransactionInfo();
        info.setXid(x);
        info.setType(XATransactionInfo.FORGET);

        try {
            // Tell the server to forget the transaction.
            this.connection.syncSendPacket(info);
        } catch (JMSException e) {
            throw toXAException(e);
        }
    }

    public boolean isSameRM(XAResource xaResource) throws XAException {
        if (xaResource == null) {
            return false;
        }
        if (!(xaResource instanceof TransactionContext)) {
            return false;
        }
        TransactionContext xar = (TransactionContext) xaResource;
        try {
            return getResourceManagerId().equals(xar.getResourceManagerId());
        } catch (Throwable e) {
            throw (XAException)new XAException("Could not get resource manager id.").initCause(e);
        }
    }


    public Xid[] recover(int flag) throws XAException {

        XATransactionInfo info = new XATransactionInfo();
        info.setType(XATransactionInfo.XA_RECOVER);

        try {
            ResponseReceipt receipt = (ResponseReceipt) this.connection.syncSendRequest(info);
            return (ActiveMQXid[]) receipt.getResult();
        } catch (JMSException e) {
            throw toXAException(e);
        }
    }

    public int getTransactionTimeout() throws XAException {

        XATransactionInfo info = new XATransactionInfo();
        info.setType(XATransactionInfo.GET_TX_TIMEOUT);

        try {
            // get the tx timeout that was set.
            IntResponseReceipt receipt = (IntResponseReceipt) this.connection.syncSendRequest(info);
            return receipt.getResult();
        } catch (JMSException e) {
            throw toXAException(e);
        }
    }

    public boolean setTransactionTimeout(int seconds) throws XAException {

        XATransactionInfo info = new XATransactionInfo();
        info.setType(XATransactionInfo.SET_TX_TIMEOUT);
        info.setTransactionTimeout(seconds);

        try {
            // Setup the new tx timeout
            this.connection.asyncSendPacket(info);
            return true;
        } catch (JMSException e) {
            throw toXAException(e);
        }
    }
    
    /////////////////////////////////////////////////////////////
    //
    // Helper methods.
    //
    /////////////////////////////////////////////////////////////
    private String getResourceManagerId() throws JMSException {
        return this.connection.getResourceManagerId();
    }
    
    private void setXid(Xid xid) {
        if (xid != null) {
            // associate
            associatedXid = xid;
            activeXid = new ActiveMQXid(xid);
        } else {
            // dis-associate
            associatedXid = null;
            activeXid = null;
        }
    }

    /**
     * Converts a JMSException from the server to an XAException.
     * if the JMSException contained a linked XAException that is
     * returned instead.
     *
     * @param e
     * @return
     */
    private XAException toXAException(JMSException e) {
        if (e.getCause() != null && e.getCause() instanceof XAException) {
            XAException original = (XAException) e.getCause();
            XAException xae = new XAException(original.getMessage());
            xae.errorCode = original.errorCode;
            xae.initCause(original);
            return xae;
        }

        XAException xae = new XAException(e.getMessage());
        xae.errorCode = XAException.XAER_RMFAIL;
        xae.initCause(e);
        return xae;
    }

    public ActiveMQConnection getConnection() {
        return connection;
    }

}