/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.apache.activemq.bugs;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.concurrent.CopyOnWriteArrayList;

import jakarta.jms.Message;
import jakarta.jms.MessageProducer;
import jakarta.jms.XASession;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;

import org.apache.activemq.ActiveMQXAConnection;
import org.apache.activemq.ActiveMQXAConnectionFactory;
import org.apache.activemq.broker.BrokerPlugin;
import org.apache.activemq.broker.BrokerPluginSupport;
import org.apache.activemq.broker.BrokerRegistry;
import org.apache.activemq.broker.BrokerRestartTestSupport;
import org.apache.activemq.broker.BrokerService;
import org.apache.activemq.broker.ConnectionContext;
import org.apache.activemq.broker.TransactionBroker;
import org.apache.activemq.broker.TransportConnection;
import org.apache.activemq.command.ConnectionId;
import org.apache.activemq.command.TransactionId;
import org.apache.activemq.command.TransactionInfo;
import org.apache.activemq.command.XATransactionId;
import org.apache.activemq.transport.failover.FailoverTransport;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * Test for AMQ-4950.
 * Simulates an error during XA prepare call.
 */
public class AMQ4950Test extends BrokerRestartTestSupport {

    protected static final Logger LOG = LoggerFactory.getLogger(AMQ4950Test.class);
    protected static final String simulatedExceptionMessage = "Simulating error inside tx prepare().";
    public boolean prioritySupport = false;
    protected String connectionUri = null;

    @Override
    protected void configureBroker(BrokerService broker) throws Exception {
        broker.setDestinationPolicy(policyMap);
        broker.setDeleteAllMessagesOnStartup(true);
        broker.setUseJmx(false);
        connectionUri = broker.addConnector("tcp://localhost:0").getPublishableConnectString();
        broker.setPlugins(new BrokerPlugin[]{
                new BrokerPluginSupport() {

                    @Override
                    public int prepareTransaction(ConnectionContext context, TransactionId xid) throws Exception {
                        getNext().prepareTransaction(context, xid);
                        LOG.debug("BrokerPlugin.prepareTransaction() will throw an exception.");
                        throw new XAException(simulatedExceptionMessage);
                    }

                    @Override
                    public void commitTransaction(ConnectionContext context, TransactionId xid, boolean onePhase) throws Exception {
                        LOG.debug("BrokerPlugin.commitTransaction().");
                        super.commitTransaction(context, xid, onePhase);
                    }
                }
        });
   }

    /**
     * Creates XA transaction and invokes XA prepare().
     * Due to registered BrokerFilter prepare will be handled by broker
     * but then throw an exception.
     * Prior to fixing AMQ-4950, this resulted in a ClassCastException
     * in ConnectionStateTracker.PrepareReadonlyTransactionAction.onResponse()
     * causing the failover transport to reconnect and replay the XA prepare().
     */
    public void testXAPrepareFailure() throws Exception {

        assertNotNull(connectionUri);
        ActiveMQXAConnectionFactory cf = new ActiveMQXAConnectionFactory("failover:(" + connectionUri + ")");
        ActiveMQXAConnection xaConnection = (ActiveMQXAConnection)cf.createConnection();
        xaConnection.start();
        XASession session = xaConnection.createXASession();
        XAResource resource = session.getXAResource();
        Xid tid = createXid();
        resource.start(tid, XAResource.TMNOFLAGS);

        MessageProducer producer = session.createProducer(session.createQueue(this.getClass().getName()));
        Message message = session.createTextMessage("Sample Message");
        producer.send(message);
        resource.end(tid, XAResource.TMSUCCESS);
        try {
            LOG.debug("Calling XA prepare(), expecting an exception");
            int ret = resource.prepare(tid);
            if (XAResource.XA_OK == ret) 
                resource.commit(tid, false);
        } catch (XAException xae) {
            LOG.info("Received excpected XAException: {}", xae.getMessage());
            LOG.info("Rolling back transaction {}", tid);
            
            // with bug AMQ-4950 the thrown error reads "Cannot call prepare now"
            // we check that we receive the original exception message as 
            // thrown by the BrokerPlugin
            assertEquals(simulatedExceptionMessage, xae.getMessage());
            resource.rollback(tid);
        }
        // couple of assertions
        assertTransactionGoneFromBroker(tid);
        assertTransactionGoneFromConnection(broker.getBrokerName(), xaConnection.getClientID(), xaConnection.getConnectionInfo().getConnectionId(), tid);
        assertTransactionGoneFromFailoverState(xaConnection, tid);

        //cleanup
        producer.close();
        session.close();
        xaConnection.close();
        LOG.debug("testXAPrepareFailure() finished.");
    }


    public Xid createXid() throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        DataOutputStream os = new DataOutputStream(baos);
        os.writeLong(++txGenerator);
        os.close();
        final byte[] bs = baos.toByteArray();

        return new Xid() {
            public int getFormatId() {
                return 86;
            }

            public byte[] getGlobalTransactionId() {
                return bs;
            }

            public byte[] getBranchQualifier() {
                return bs;
            }
        };
    }


    private void assertTransactionGoneFromFailoverState(
            ActiveMQXAConnection connection1, Xid tid) throws Exception {

        FailoverTransport transport = (FailoverTransport) connection1.getTransport().narrow(FailoverTransport.class);
        TransactionInfo info = new TransactionInfo(connection1.getConnectionInfo().getConnectionId(), new XATransactionId(tid), TransactionInfo.COMMIT_ONE_PHASE);
        assertNull("transaction should not exist in the state tracker",
                transport.getStateTracker().processCommitTransactionOnePhase(info));
    }


    private void assertTransactionGoneFromBroker(Xid tid) throws Exception {
        BrokerService broker = BrokerRegistry.getInstance().lookup("localhost");
        TransactionBroker transactionBroker = (TransactionBroker)broker.getBroker().getAdaptor(TransactionBroker.class);
        try {
            transactionBroker.getTransaction(null, new XATransactionId(tid), false);
            fail("expected exception on tx not found");
        } catch (XAException expectedOnNotFound) {
        }
    }


    private void assertTransactionGoneFromConnection(String brokerName, String clientId, ConnectionId connectionId, Xid tid) throws Exception {
        BrokerService broker = BrokerRegistry.getInstance().lookup(brokerName);
        CopyOnWriteArrayList<TransportConnection> connections = broker.getTransportConnectors().get(0).getConnections();
        for (TransportConnection connection: connections) {
            if (connection.getConnectionId().equals(clientId)) {
                try {
                    connection.processPrepareTransaction(new TransactionInfo(connectionId, new XATransactionId(tid), TransactionInfo.PREPARE));
                    fail("did not get expected excepton on missing transaction, it must be still there in error!");
                } catch (IllegalStateException expectedOnNoTransaction) {
                }
            }
        }
    }
}
