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

import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;

import org.apache.activemq.broker.ConnectionContext;
import org.apache.activemq.command.ActiveMQDestination;
import org.apache.activemq.command.ActiveMQQueue;
import org.apache.activemq.command.ActiveMQTopic;
import org.apache.activemq.command.Message;
import org.apache.activemq.command.MessageAck;
import org.apache.activemq.command.MessageId;
import org.apache.activemq.command.TransactionId;
import org.apache.activemq.command.XATransactionId;
import org.apache.activemq.store.IndexListener;
import org.apache.activemq.store.MessageStore;
import org.apache.activemq.store.ProxyMessageStore;
import org.apache.activemq.store.TopicMessageStore;
import org.apache.activemq.store.TransactionRecoveryListener;
import org.apache.activemq.store.memory.MemoryTransactionStore;
import org.apache.activemq.util.ByteSequence;
import org.apache.activemq.util.DataByteArrayInputStream;

/**
 * respect 2pc prepare
 * uses local transactions to maintain prepared state
 * xid column provides transaction flag for additions and removals
 * a commit clears that context and completes the work
 * a rollback clears the flag and removes the additions
 * Essentially a prepare is an insert &| update transaction
 *  commit|rollback is an update &| remove
 */
public class JdbcMemoryTransactionStore extends MemoryTransactionStore {


    public JdbcMemoryTransactionStore(JDBCPersistenceAdapter jdbcPersistenceAdapter) {
        super(jdbcPersistenceAdapter, jdbcPersistenceAdapter.getBrokerService());
    }

    @Override
    public void prepare(TransactionId txid) throws IOException {
        Tx tx = inflightTransactions.remove(txid);
        if (tx == null) {
            return;
        }

        ConnectionContext ctx = new ConnectionContext();
        // setting the xid modifies the add/remove to be pending transaction outcome
        ctx.setXid((XATransactionId) txid);
        persistenceAdapter.beginTransaction(ctx);
        try {

            // Do all the message adds.
            for (Iterator<AddMessageCommand> iter = tx.messages.iterator(); iter.hasNext();) {
                AddMessageCommand cmd = iter.next();
                cmd.run(ctx);
            }
            // And removes..
            for (Iterator<RemoveMessageCommand> iter = tx.acks.iterator(); iter.hasNext();) {
                RemoveMessageCommand cmd = iter.next();
                cmd.run(ctx);
            }

            persistenceAdapter.commitTransaction(ctx);

        } catch ( IOException e ) {
            persistenceAdapter.rollbackTransaction(ctx);
            throw e;
        }

        ctx.setXid(null);
        // setup for commit outcome
        ArrayList<AddMessageCommand> updateFromPreparedStateCommands = new ArrayList<AddMessageCommand>();
        for (Iterator<AddMessageCommand> iter = tx.messages.iterator(); iter.hasNext();) {
            final AddMessageCommand addMessageCommand = iter.next();
            updateFromPreparedStateCommands.add(new CommitAddOutcome(addMessageCommand));
        }
        tx.messages = updateFromPreparedStateCommands;
        preparedTransactions.put(txid, tx);

    }


    class CommitAddOutcome implements AddMessageCommand {
        final Message message;
        JDBCMessageStore jdbcMessageStore;

        public CommitAddOutcome(JDBCMessageStore jdbcMessageStore, Message message) {
            this.jdbcMessageStore = jdbcMessageStore;
            this.message = message;
        }

        public CommitAddOutcome(AddMessageCommand addMessageCommand) {
            this((JDBCMessageStore)addMessageCommand.getMessageStore(), addMessageCommand.getMessage());
        }

        @Override
        public Message getMessage() {
            return message;
        }

        @Override
        public MessageStore getMessageStore() {
            return jdbcMessageStore;
        }

        @Override
        public void run(final ConnectionContext context) throws IOException {
            JDBCPersistenceAdapter jdbcPersistenceAdapter = (JDBCPersistenceAdapter) persistenceAdapter;
            final Long preparedEntrySequence = (Long) message.getMessageId().getEntryLocator();
            TransactionContext c = jdbcPersistenceAdapter.getTransactionContext(context);

            long newSequence;
            synchronized (jdbcMessageStore.pendingAdditions) {
                newSequence = jdbcPersistenceAdapter.getNextSequenceId();
                final long sequenceToSet = newSequence;
                c.onCompletion(new Runnable() {
                    @Override
                    public void run() {
                        message.getMessageId().setEntryLocator(sequenceToSet);
                        message.getMessageId().setFutureOrSequenceLong(sequenceToSet);
                    }
                });

                if (jdbcMessageStore.getIndexListener() != null) {
                    jdbcMessageStore.getIndexListener().onAdd(new IndexListener.MessageContext(context, message, null));
                }
            }

            jdbcPersistenceAdapter.commitAdd(context, message.getMessageId(), preparedEntrySequence, newSequence);
            jdbcMessageStore.onAdd(message, (Long)message.getMessageId().getEntryLocator(), message.getPriority());
        }

        @Override
        public void setMessageStore(MessageStore messageStore) {
            jdbcMessageStore = (JDBCMessageStore) messageStore;
        }
    }

    @Override
    public void rollback(TransactionId txid) throws IOException {

        Tx tx = inflightTransactions.remove(txid);
        if (tx == null) {
            tx = preparedTransactions.remove(txid);
            if (tx != null) {
                // undo prepare work
                ConnectionContext ctx = new ConnectionContext();
                try {
                    if (brokerService != null) {
                        ctx.setBroker(brokerService.getBroker());
                    }
                }  catch (Exception e) {
                    throw new IOException(e.getMessage(), e);
                }
                persistenceAdapter.beginTransaction(ctx);
                try {

                    for (Iterator<AddMessageCommand> iter = tx.messages.iterator(); iter.hasNext(); ) {
                        final Message message = iter.next().getMessage();
                        // need to delete the row
                        ((JDBCPersistenceAdapter) persistenceAdapter).commitRemove(ctx, new MessageAck(message, MessageAck.STANDARD_ACK_TYPE, 1));
                    }

                    for (Iterator<RemoveMessageCommand> iter = tx.acks.iterator(); iter.hasNext(); ) {
                        RemoveMessageCommand removeMessageCommand = iter.next();
                        if (removeMessageCommand instanceof LastAckCommand ) {
                            ((LastAckCommand)removeMessageCommand).rollback(ctx);
                        } else {
                            MessageId messageId = removeMessageCommand.getMessageAck().getLastMessageId();
                            long sequence = (Long)messageId.getEntryLocator();
                            // need to unset the txid flag on the existing row
                            ((JDBCPersistenceAdapter) persistenceAdapter).commitAdd(ctx, messageId, sequence, sequence);

                            if (removeMessageCommand instanceof RecoveredRemoveMessageCommand) {
                                ((JDBCMessageStore) removeMessageCommand.getMessageStore()).trackRollbackAck(((RecoveredRemoveMessageCommand) removeMessageCommand).getMessage());
                            }
                        }
                    }
                } catch (IOException e) {
                    persistenceAdapter.rollbackTransaction(ctx);
                    throw e;
                }
                persistenceAdapter.commitTransaction(ctx);
            }
        }
    }

    @Override
    public void recover(TransactionRecoveryListener listener) throws IOException {
        ((JDBCPersistenceAdapter)persistenceAdapter).recover(this);
        super.recover(listener);
    }

    public void recoverAdd(long id, byte[] messageBytes) throws IOException {
        final Message message = (Message) ((JDBCPersistenceAdapter)persistenceAdapter).getWireFormat().unmarshal(new ByteSequence(messageBytes));
        message.getMessageId().setFutureOrSequenceLong(id);
        message.getMessageId().setEntryLocator(id);
        Tx tx = getPreparedTx(message.getTransactionId());
        tx.add(new CommitAddOutcome(null, message));
    }

    public void recoverAck(long id, byte[] xid, byte[] message) throws IOException {
        final Message msg = (Message) ((JDBCPersistenceAdapter)persistenceAdapter).getWireFormat().unmarshal(new ByteSequence(message));
        msg.getMessageId().setFutureOrSequenceLong(id);
        msg.getMessageId().setEntryLocator(id);
        Tx tx = getPreparedTx(new XATransactionId(xid));
        final MessageAck ack = new MessageAck(msg, MessageAck.STANDARD_ACK_TYPE, 1);
        tx.add(new RecoveredRemoveMessageCommand() {
            MessageStore messageStore = null;
            @Override
            public MessageAck getMessageAck() {
                return ack;
            }

            @Override
            public void run(ConnectionContext context) throws IOException {
                ((JDBCPersistenceAdapter)persistenceAdapter).commitRemove(context, ack);
            }

            public Message getMessage() {
                return msg;
            }

            @Override
            public void setMessageStore(MessageStore messageStore) {
                this.messageStore = messageStore;
            }

            @Override
            public MessageStore getMessageStore() {
                return messageStore;
            }

        });
    }

    interface RecoveredRemoveMessageCommand extends RemoveMessageCommand {
        Message getMessage();

        void setMessageStore(MessageStore messageStore);
    }

    interface LastAckCommand extends RemoveMessageCommand {
        void rollback(ConnectionContext context) throws IOException;

        String getClientId();

        String getSubName();

        long getSequence();

        byte getPriority();

        void setMessageStore(JDBCTopicMessageStore jdbcTopicMessageStore);
    }

    public void recoverLastAck(byte[] encodedXid, final ActiveMQDestination destination, final String subName, final String clientId) throws IOException {
        Tx tx = getPreparedTx(new XATransactionId(encodedXid));
        DataByteArrayInputStream inputStream = new DataByteArrayInputStream(encodedXid);
        inputStream.skipBytes(1); // +|-
        final long lastAck = inputStream.readLong();
        final byte priority = inputStream.readByte();
        final MessageAck ack = new MessageAck();
        ack.setDestination(destination);
        tx.add(new LastAckCommand() {
            JDBCTopicMessageStore jdbcTopicMessageStore;

            @Override
            public MessageAck getMessageAck() {
                return ack;
            }

            @Override
            public MessageStore getMessageStore() {
                return jdbcTopicMessageStore;
            }

            @Override
            public void run(ConnectionContext context) throws IOException {
                ((JDBCPersistenceAdapter)persistenceAdapter).commitLastAck(context, lastAck, priority, destination, subName, clientId);
                jdbcTopicMessageStore.complete(clientId, subName);
            }

            @Override
            public void rollback(ConnectionContext context) throws IOException {
                ((JDBCPersistenceAdapter)persistenceAdapter).rollbackLastAck(context, jdbcTopicMessageStore.isPrioritizedMessages() ? priority : 0, jdbcTopicMessageStore.getDestination(), subName, clientId);
                jdbcTopicMessageStore.complete(clientId, subName);
            }

            @Override
            public String getClientId() {
                return clientId;
            }

            @Override
            public String getSubName() {
                return subName;
            }

            @Override
            public long getSequence() {
                return lastAck;
            }

            @Override
            public byte getPriority() {
                return priority;
            }

            @Override
            public void setMessageStore(JDBCTopicMessageStore jdbcTopicMessageStore) {
                this.jdbcTopicMessageStore = jdbcTopicMessageStore;
            }
        });

    }

    @Override
    protected void onRecovered(Tx tx) {
        for (RemoveMessageCommand removeMessageCommand: tx.acks) {
            if (removeMessageCommand instanceof LastAckCommand) {
                LastAckCommand lastAckCommand = (LastAckCommand) removeMessageCommand;
                JDBCTopicMessageStore jdbcTopicMessageStore = (JDBCTopicMessageStore) findMessageStore(lastAckCommand.getMessageAck().getDestination());
                jdbcTopicMessageStore.pendingCompletion(lastAckCommand.getClientId(), lastAckCommand.getSubName(), lastAckCommand.getSequence(), lastAckCommand.getPriority());
                lastAckCommand.setMessageStore(jdbcTopicMessageStore);
            } else {
                ((RecoveredRemoveMessageCommand)removeMessageCommand).setMessageStore(findMessageStore(removeMessageCommand.getMessageAck().getDestination()));
            }
        }
        for (AddMessageCommand addMessageCommand : tx.messages) {
            addMessageCommand.setMessageStore(findMessageStore(addMessageCommand.getMessage().getDestination()));
        }
    }

    private MessageStore findMessageStore(ActiveMQDestination destination) {
        ProxyMessageStore proxyMessageStore = null;
        try {
            if (destination.isQueue()) {
                proxyMessageStore = (ProxyMessageStore) persistenceAdapter.createQueueMessageStore((ActiveMQQueue) destination);
            } else {
                proxyMessageStore = (ProxyMessageStore) persistenceAdapter.createTopicMessageStore((ActiveMQTopic) destination);
            }
        } catch (IOException error) {
            throw new RuntimeException("Failed to find/create message store for destination: " + destination, error);
        }
        return proxyMessageStore.getDelegate();
    }

    @Override
    public void acknowledge(final TopicMessageStore topicMessageStore, final String clientId, final String subscriptionName,
                           final MessageId messageId, final MessageAck ack) throws IOException {

        if (ack.isInTransaction()) {
            Tx tx = getTx(ack.getTransactionId());
            tx.add(new LastAckCommand() {
                public MessageAck getMessageAck() {
                    return ack;
                }

                public void run(ConnectionContext ctx) throws IOException {
                    topicMessageStore.acknowledge(ctx, clientId, subscriptionName, messageId, ack);
                }

                @Override
                public MessageStore getMessageStore() {
                    return topicMessageStore;
                }

                @Override
                public void rollback(ConnectionContext context) throws IOException {
                    JDBCTopicMessageStore jdbcTopicMessageStore = (JDBCTopicMessageStore)topicMessageStore;
                    ((JDBCPersistenceAdapter)persistenceAdapter).rollbackLastAck(context,
                            jdbcTopicMessageStore,
                            ack,
                            subscriptionName, clientId);
                    jdbcTopicMessageStore.complete(clientId, subscriptionName);
                }


                @Override
                public String getClientId() {
                    return clientId;
                }

                @Override
                public String getSubName() {
                    return subscriptionName;
                }

                @Override
                public long getSequence() {
                    throw new IllegalStateException("Sequence id must be inferred from ack");
                }

                @Override
                public byte getPriority() {
                    throw new IllegalStateException("Priority must be inferred from ack or row");
                }

                @Override
                public void setMessageStore(JDBCTopicMessageStore jdbcTopicMessageStore) {
                    throw new IllegalStateException("message store already known!");
                }
            });
        } else {
            topicMessageStore.acknowledge(null, clientId, subscriptionName, messageId, ack);
        }
    }

}
