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

import java.io.File;
import java.io.IOException;
import java.util.Date;
import java.util.HashSet;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.activemq.broker.Broker;
import org.apache.activemq.broker.ConnectionContext;
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.AbstractMessageStore;
import org.apache.activemq.store.ListenableFuture;
import org.apache.activemq.store.MessageStore;
import org.apache.activemq.store.PersistenceAdapter;
import org.apache.activemq.store.ProxyMessageStore;
import org.apache.activemq.store.ProxyTopicMessageStore;
import org.apache.activemq.store.TopicMessageStore;
import org.apache.activemq.store.TransactionRecoveryListener;
import org.apache.activemq.store.TransactionStore;
import org.apache.activemq.store.kahadb.data.KahaCommitCommand;
import org.apache.activemq.store.kahadb.data.KahaEntryType;
import org.apache.activemq.store.kahadb.data.KahaPrepareCommand;
import org.apache.activemq.store.kahadb.data.KahaTraceCommand;
import org.apache.activemq.store.kahadb.disk.journal.Journal;
import org.apache.activemq.store.kahadb.disk.journal.Location;
import org.apache.activemq.util.DataByteArrayInputStream;
import org.apache.activemq.util.DataByteArrayOutputStream;
import org.apache.activemq.util.IOHelper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class MultiKahaDBTransactionStore implements TransactionStore {
    static final Logger LOG = LoggerFactory.getLogger(MultiKahaDBTransactionStore.class);
    final MultiKahaDBPersistenceAdapter multiKahaDBPersistenceAdapter;
    final ConcurrentMap<TransactionId, Tx> inflightTransactions = new ConcurrentHashMap<TransactionId, Tx>();
    final Set<TransactionId> recoveredPendingCommit = new HashSet<TransactionId>();
    private Journal journal;
    private int journalMaxFileLength = Journal.DEFAULT_MAX_FILE_LENGTH;
    private int journalWriteBatchSize = Journal.DEFAULT_MAX_WRITE_BATCH_SIZE;
    private final AtomicBoolean started = new AtomicBoolean(false);

    public MultiKahaDBTransactionStore(MultiKahaDBPersistenceAdapter multiKahaDBPersistenceAdapter) {
        this.multiKahaDBPersistenceAdapter = multiKahaDBPersistenceAdapter;
    }

    public MessageStore proxy(final TransactionStore transactionStore, MessageStore messageStore) {
        return new ProxyMessageStore(messageStore) {
            @Override
            public void addMessage(ConnectionContext context, final Message send) throws IOException {
                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
            }

            @Override
            public void addMessage(ConnectionContext context, final Message send, boolean canOptimizeHint) throws IOException {
                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
            }

            @Override
            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message) throws IOException {
                return MultiKahaDBTransactionStore.this.asyncAddQueueMessage(transactionStore, context, getDelegate(), message);
            }

            @Override
            public ListenableFuture<Object> asyncAddQueueMessage(ConnectionContext context, Message message, boolean canOptimizeHint) throws IOException {
                return MultiKahaDBTransactionStore.this.asyncAddQueueMessage(transactionStore, context, getDelegate(), message);
            }

            @Override
            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
                MultiKahaDBTransactionStore.this.removeMessage(transactionStore, context, getDelegate(), ack);
            }

            @Override
            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
                MultiKahaDBTransactionStore.this.removeAsyncMessage(transactionStore, context, getDelegate(), ack);
            }
        };
    }

    public TopicMessageStore proxy(final TransactionStore transactionStore, final TopicMessageStore messageStore) {
        return new ProxyTopicMessageStore(messageStore) {
            @Override
            public void addMessage(ConnectionContext context, final Message send, boolean canOptimizeHint) throws IOException {
                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
            }

            @Override
            public void addMessage(ConnectionContext context, final Message send) throws IOException {
                MultiKahaDBTransactionStore.this.addMessage(transactionStore, context, getDelegate(), send);
            }

            @Override
            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message, boolean canOptimizeHint) throws IOException {
                return MultiKahaDBTransactionStore.this.asyncAddTopicMessage(transactionStore, context, getDelegate(), message);
            }

            @Override
            public ListenableFuture<Object> asyncAddTopicMessage(ConnectionContext context, Message message) throws IOException {
                return MultiKahaDBTransactionStore.this.asyncAddTopicMessage(transactionStore, context, getDelegate(), message);
            }

            @Override
            public void removeMessage(ConnectionContext context, final MessageAck ack) throws IOException {
                MultiKahaDBTransactionStore.this.removeMessage(transactionStore, context, getDelegate(), ack);
            }

            @Override
            public void removeAsyncMessage(ConnectionContext context, MessageAck ack) throws IOException {
                MultiKahaDBTransactionStore.this.removeAsyncMessage(transactionStore, context, getDelegate(), ack);
            }

            @Override
            public void acknowledge(ConnectionContext context, String clientId, String subscriptionName,
                                    MessageId messageId, MessageAck ack) throws IOException {
                MultiKahaDBTransactionStore.this.acknowledge(transactionStore, context, (TopicMessageStore) getDelegate(), clientId,
                        subscriptionName, messageId, ack);
            }
        };
    }

    public void deleteAllMessages() {
        IOHelper.deleteChildren(getDirectory());
    }

    public int getJournalMaxFileLength() {
        return journalMaxFileLength;
    }

    public void setJournalMaxFileLength(int journalMaxFileLength) {
        this.journalMaxFileLength = journalMaxFileLength;
    }

    public int getJournalMaxWriteBatchSize() {
        return journalWriteBatchSize;
    }

    public void setJournalMaxWriteBatchSize(int journalWriteBatchSize) {
        this.journalWriteBatchSize = journalWriteBatchSize;
    }

    public class Tx {
        private final Set<TransactionStore> stores = new HashSet<TransactionStore>();
        private int prepareLocationId = 0;

        public void trackStore(TransactionStore store) {
            stores.add(store);
        }

        public Set<TransactionStore> getStores() {
            return stores;
        }

        public void trackPrepareLocation(Location location) {
            this.prepareLocationId = location.getDataFileId();
        }

        public int getPreparedLocationId() {
            return prepareLocationId;
        }
    }

    public Tx getTx(TransactionId txid) {
        Tx tx = inflightTransactions.get(txid);
        if (tx == null) {
            tx = new Tx();
            inflightTransactions.put(txid, tx);
        }
        return tx;
    }

    public Tx removeTx(TransactionId txid) {
        return inflightTransactions.remove(txid);
    }

    @Override
    public void prepare(TransactionId txid) throws IOException {
        Tx tx = getTx(txid);
        for (TransactionStore store : tx.getStores()) {
            store.prepare(txid);
        }
    }

    @Override
    public void commit(TransactionId txid, boolean wasPrepared, Runnable preCommit, Runnable postCommit)
            throws IOException {

        if (preCommit != null) {
            preCommit.run();
        }

        Tx tx = getTx(txid);
        if (wasPrepared) {
            for (TransactionStore store : tx.getStores()) {
                store.commit(txid, true, null, null);
            }
        } else {
            // can only do 1pc on a single store
            if (tx.getStores().size() == 1) {
                for (TransactionStore store : tx.getStores()) {
                    store.commit(txid, false, null, null);
                }
            } else {
                // need to do local 2pc
                for (TransactionStore store : tx.getStores()) {
                    store.prepare(txid);
                }
                persistOutcome(tx, txid);
                for (TransactionStore store : tx.getStores()) {
                    store.commit(txid, true, null, null);
                }
                persistCompletion(txid);
            }
        }
        removeTx(txid);
        if (postCommit != null) {
            postCommit.run();
        }
    }

    public void persistOutcome(Tx tx, TransactionId txid) throws IOException {
        tx.trackPrepareLocation(store(new KahaPrepareCommand().setTransactionInfo(TransactionIdConversion.convert(multiKahaDBPersistenceAdapter.transactionIdTransformer.transform(txid)))));
    }

    public void persistCompletion(TransactionId txid) throws IOException {
        store(new KahaCommitCommand().setTransactionInfo(TransactionIdConversion.convert(multiKahaDBPersistenceAdapter.transactionIdTransformer.transform(txid))));
    }

    private Location store(JournalCommand<?> data) throws IOException {
        int size = data.serializedSizeFramed();
        DataByteArrayOutputStream os = new DataByteArrayOutputStream(size + 1);
        os.writeByte(data.type().getNumber());
        data.writeFramed(os);
        Location location = journal.write(os.toByteSequence(), true);
        journal.setLastAppendLocation(location);
        return location;
    }

    @Override
    public void rollback(TransactionId txid) throws IOException {
        Tx tx = removeTx(txid);
        if (tx != null) {
            for (TransactionStore store : tx.getStores()) {
                store.rollback(txid);
            }
        }
    }

    @Override
    public void start() throws Exception {
        if (started.compareAndSet(false, true)) {
            journal = new Journal() {
                @Override
                public void cleanup() {
                    super.cleanup();
                    txStoreCleanup();
                }
            };
            journal.setDirectory(getDirectory());
            journal.setMaxFileLength(journalMaxFileLength);
            journal.setWriteBatchSize(journalWriteBatchSize);
            IOHelper.mkdirs(journal.getDirectory());
            journal.start();
            recoverPendingLocalTransactions();
            store(new KahaTraceCommand().setMessage("LOADED " + new Date()));
        }
    }

    private void txStoreCleanup() {
        Set<Integer> knownDataFileIds = new TreeSet<Integer>(journal.getFileMap().keySet());
        for (Tx tx : inflightTransactions.values()) {
            knownDataFileIds.remove(tx.getPreparedLocationId());
        }
        try {
            journal.removeDataFiles(knownDataFileIds);
        } catch (Exception e) {
            LOG.error(this + ", Failed to remove tx journal datafiles " + knownDataFileIds);
        }
    }

    private File getDirectory() {
        return new File(multiKahaDBPersistenceAdapter.getDirectory(), "txStore");
    }

    @Override
    public void stop() throws Exception {
        if (started.compareAndSet(true, false) && journal != null) {
            journal.close();
            journal = null;
        }
    }

    private void recoverPendingLocalTransactions() throws IOException {
        Location location = journal.getNextLocation(null);
        while (location != null) {
            process(load(location));
            location = journal.getNextLocation(location);
        }
        recoveredPendingCommit.addAll(inflightTransactions.keySet());
        LOG.info("pending local transactions: " + recoveredPendingCommit);
    }

    public JournalCommand<?> load(Location location) throws IOException {
        DataByteArrayInputStream is = new DataByteArrayInputStream(journal.read(location));
        byte readByte = is.readByte();
        KahaEntryType type = KahaEntryType.valueOf(readByte);
        if (type == null) {
            throw new IOException("Could not load journal record. Invalid location: " + location);
        }
        JournalCommand<?> message = (JournalCommand<?>) type.createMessage();
        message.mergeFramed(is);
        return message;
    }

    public void process(JournalCommand<?> command) throws IOException {
        switch (command.type()) {
            case KAHA_PREPARE_COMMAND:
                KahaPrepareCommand prepareCommand = (KahaPrepareCommand) command;
                getTx(TransactionIdConversion.convert(prepareCommand.getTransactionInfo()));
                break;
            case KAHA_COMMIT_COMMAND:
                KahaCommitCommand commitCommand = (KahaCommitCommand) command;
                removeTx(TransactionIdConversion.convert(commitCommand.getTransactionInfo()));
                break;
            case KAHA_TRACE_COMMAND:
                break;
            default:
                throw new IOException("Unexpected command in transaction journal: " + command);
        }
    }


    @Override
    public synchronized void recover(final TransactionRecoveryListener listener) throws IOException {

        for (final PersistenceAdapter adapter : multiKahaDBPersistenceAdapter.adapters) {
            adapter.createTransactionStore().recover(new TransactionRecoveryListener() {
                @Override
                public void recover(XATransactionId xid, Message[] addedMessages, MessageAck[] acks) {
                    try {
                        getTx(xid).trackStore(adapter.createTransactionStore());
                    } catch (IOException e) {
                        LOG.error("Failed to access transaction store: " + adapter + " for prepared xa tid: " + xid, e);
                    }
                    listener.recover(xid, addedMessages, acks);
                }
            });
        }

        try {
            Broker broker = multiKahaDBPersistenceAdapter.getBrokerService().getBroker();
            // force completion of local xa
            for (TransactionId txid : broker.getPreparedTransactions(null)) {
                if (multiKahaDBPersistenceAdapter.isLocalXid(txid)) {
                    try {
                        if (recoveredPendingCommit.contains(txid)) {
                            LOG.info("delivering pending commit outcome for tid: " + txid);
                            broker.commitTransaction(null, txid, false);

                        } else {
                            LOG.info("delivering rollback outcome to store for tid: " + txid);
                            broker.forgetTransaction(null, txid);
                        }
                        persistCompletion(txid);
                    } catch (Exception ex) {
                        LOG.error("failed to deliver pending outcome for tid: " + txid, ex);
                    }
                }
            }
        } catch (Exception e) {
            LOG.error("failed to resolve pending local transactions", e);
        }
    }

    void addMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
            throws IOException {
        if (message.getTransactionId() != null) {
            getTx(message.getTransactionId()).trackStore(transactionStore);
        }
        destination.addMessage(context, message);
    }

    ListenableFuture<Object> asyncAddQueueMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
            throws IOException {
        if (message.getTransactionId() != null) {
            getTx(message.getTransactionId()).trackStore(transactionStore);
            destination.addMessage(context, message);
            return AbstractMessageStore.FUTURE;
        } else {
            return destination.asyncAddQueueMessage(context, message);
        }
    }

    ListenableFuture<Object> asyncAddTopicMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final Message message)
            throws IOException {

        if (message.getTransactionId() != null) {
            getTx(message.getTransactionId()).trackStore(transactionStore);
            destination.addMessage(context, message);
            return AbstractMessageStore.FUTURE;
        } else {
            return destination.asyncAddTopicMessage(context, message);
        }
    }

    final void removeMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final MessageAck ack)
            throws IOException {
        if (ack.getTransactionId() != null) {
            getTx(ack.getTransactionId()).trackStore(transactionStore);
        }
        destination.removeMessage(context, ack);
    }

    final void removeAsyncMessage(final TransactionStore transactionStore, ConnectionContext context, final MessageStore destination, final MessageAck ack)
            throws IOException {
        if (ack.getTransactionId() != null) {
            getTx(ack.getTransactionId()).trackStore(transactionStore);
        }
        destination.removeAsyncMessage(context, ack);
    }

    final void acknowledge(final TransactionStore transactionStore, ConnectionContext context, final TopicMessageStore destination,
                           final String clientId, final String subscriptionName,
                           final MessageId messageId, final MessageAck ack) throws IOException {
        if (ack.getTransactionId() != null) {
            getTx(ack.getTransactionId()).trackStore(transactionStore);
        }
        destination.acknowledge(context, clientId, subscriptionName, messageId, ack);
    }

}
