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

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.codehaus.activemq.message.ActiveMQXid;
import org.codehaus.activemq.message.IntResponseReceipt;
import org.codehaus.activemq.message.ResponseReceipt;
import org.codehaus.activemq.message.XATransactionInfo;

import javax.jms.*;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;

/**
 * The XASession interface extends the capability of Session by adding access
 * to a JMS provider's support for the  Java Transaction API (JTA) (optional).
 * This support takes the form of a javax.transaction.xa.XAResource object.
 * The functionality of this object closely resembles that defined by the
 * standard X/Open XA Resource 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/>
 * A client of the application server is given what it thinks is a regular
 * JMS Session. Behind the scenes, the application server controls the
 * transaction management of the underlying XASession.
 * <p/>
 * The XASession interface is optional. JMS providers are not required to
 * support this interface. This interface is for use by JMS providers to
 * support transactional environments. Client programs are strongly encouraged
 * to use the transactional support  available in their environment, rather
 * than use these XA  interfaces directly.
 *
 * @version $Revision: 1.3 $
 * @see javax.jms.Session
 * @see javax.jms.QueueSession
 * @see javax.jms.TopicSession
 * @see javax.jms.XASession
 */
public class ActiveMQXASession extends ActiveMQSession implements QueueSession, TopicSession, XAQueueSession, XATopicSession, XAResource {
    private static final Log log = LogFactory.getLog(ActiveMQXASession.class);
    private Xid associatedXid;
    private ActiveMQXid activeXid;

    public ActiveMQXASession(ActiveMQXAConnection theConnection, int theAcknowlegeMode) throws JMSException {
        super(theConnection, theAcknowlegeMode);
    }

    public boolean getTransacted() throws JMSException {
        return true;
    }

    public void rollback() throws JMSException {
        throw new TransactionInProgressException("Cannot rollback() inside an XASession");
    }

    public void commit() throws JMSException {
        throw new TransactionInProgressException("Cannot commit() inside an XASession");
    }

    public Session getSession() throws JMSException {
        return this;
    }

    public XAResource getXAResource() {
        return this;
    }

    public QueueSession getQueueSession() throws JMSException {
        return this;
    }

    public TopicSession getTopicSession() throws JMSException {
        return this;
    }

    /**
     * Associates a transaction with the resource.
     */
    public void start(Xid xid, int flags) throws XAException {
        checkClosedXA();

        // 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.setId(this.packetIdGenerator.generateId());
        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 {
        checkClosedXA();

        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)) {
//            x = activeXid;
            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.setId(this.packetIdGenerator.generateId());
        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.setId(this.packetIdGenerator.generateId());
        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);
        }
    }

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

        // 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
//            x = activeXid;
            throw new XAException(XAException.XAER_PROTO);
        } else {
            x = new ActiveMQXid(xid);
        }

        XATransactionInfo info = new XATransactionInfo();
        info.setId(this.packetIdGenerator.generateId());
        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);
        }
    }


    public void forget(Xid xid) throws XAException {
        checkClosedXA();

        // 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.setId(this.packetIdGenerator.generateId());
        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);
        }
    }

    private String getResourceManagerId() {
        return ((ActiveMQXAConnection) this.connection).getResourceManagerId();
    }

    public boolean isSameRM(XAResource xaResource) throws XAException {
        if (xaResource == null) {
            return false;
        }
        if (!(xaResource instanceof ActiveMQXASession)) {
            return false;
        }
        ActiveMQXASession xar = (ActiveMQXASession) xaResource;
        return getResourceManagerId().equals(xar.getResourceManagerId());
    }


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

        XATransactionInfo info = new XATransactionInfo();
        info.setId(this.packetIdGenerator.generateId());
        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 {
        checkClosedXA();

        XATransactionInfo info = new XATransactionInfo();
        info.setId(this.packetIdGenerator.generateId());
        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 {
        checkClosedXA();

        XATransactionInfo info = new XATransactionInfo();
        info.setId(this.packetIdGenerator.generateId());
        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);
        }
    }
    
    /**
     * overide Session - which needs to rollback if transacted
     */
    public void close() throws JMSException {
        if (!this.closed.get()) {
            doClose();
            closed.set(true);
        }      
    }


    /**
     * @throws XAException if the Session is closed
     */
    protected void checkClosedXA() throws XAException {
        if (this.closed.get()) {
            throw new XAException(XAException.XAER_RMFAIL);
        }
    }

    protected boolean isXaTransacted() {
        return true;
    }

    /**
     * This is called before transacted work is done by
     * the session.  XA Work can only be done when this
     * XA resource is associated with an Xid.
     *
     * @throws JMSException not associated with an Xid
     */
    protected void doStartTransaction() throws JMSException {
        if (associatedXid == null) {
            throw new JMSException("Session's XAResource has not been enlisted in a distributed transaction.");
        }
    }


    private void setXid(Xid xid) {
        if (xid != null) {
// associate
            associatedXid = xid;
            activeXid = new ActiveMQXid(xid);
            super.currentTransactionId = activeXid.toLocalTransactionId();
        } else {
// dis-associate
            associatedXid = null;
            activeXid = null;
            super.currentTransactionId = 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;
    }

}