/*
 * 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.bookkeeper.statelib.impl.journal;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;
import java.io.IOException;
import java.time.Duration;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.common.concurrent.FutureEventListener;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.statelib.api.AsyncStateStore;
import org.apache.bookkeeper.statelib.api.StateStore;
import org.apache.bookkeeper.statelib.api.StateStoreSpec;
import org.apache.bookkeeper.statelib.api.exceptions.InvalidStateStoreException;
import org.apache.bookkeeper.statelib.api.exceptions.StateStoreClosedException;
import org.apache.bookkeeper.statelib.api.exceptions.StateStoreException;
import org.apache.bookkeeper.statelib.api.exceptions.StateStoreRuntimeException;
import org.apache.bookkeeper.statelib.impl.Constants;
import org.apache.distributedlog.DLSN;
import org.apache.distributedlog.LogRecord;
import org.apache.distributedlog.LogRecordWithDLSN;
import org.apache.distributedlog.api.AsyncLogReader;
import org.apache.distributedlog.api.AsyncLogWriter;
import org.apache.distributedlog.api.DistributedLogManager;
import org.apache.distributedlog.api.namespace.Namespace;
import org.apache.distributedlog.bk.LedgerMetadata;
import org.apache.distributedlog.exceptions.LogEmptyException;
import org.apache.distributedlog.exceptions.LogNotFoundException;
import org.apache.distributedlog.util.Utils;

/**
 * An abstract implementation of {@link AsyncStateStore} with journal.
 */
@Slf4j
public abstract class AbstractStateStoreWithJournal<LocalStateStoreT extends StateStore> implements AsyncStateStore {

    // local state store instance
    @Getter
    protected final LocalStateStoreT localStore;

    // parameters
    protected String name = "UNINITIALIZED";

    // spec
    protected StateStoreSpec spec;

    protected boolean ownWriteScheduler = false;
    protected boolean ownReadScheduler = false;
    protected ScheduledExecutorService writeIOScheduler;
    protected ScheduledExecutorService readIOScheduler;

    // journal
    private final Namespace logNamespace;
    private DistributedLogManager logManager;
    private AsyncLogWriter writer;
    private long nextRevision;
    private CommandProcessor<LocalStateStoreT> commandProcessor;

    // checkpoint
    private ScheduledFuture<?> checkpointTask;
    private Duration checkpointInterval;

    // close state
    protected boolean isInitialized = false;
    protected CompletableFuture<Void> closeFuture = null;

    protected AbstractStateStoreWithJournal(Supplier<LocalStateStoreT> localStateStoreSupplier,
                                            Supplier<Namespace> namespaceSupplier) {
        this.localStore = localStateStoreSupplier.get();
        this.logNamespace = namespaceSupplier.get();
    }

    protected boolean ownWriteScheduler() {
        return ownWriteScheduler;
    }

    protected boolean ownReadScheduler() {
        return ownReadScheduler;
    }

    synchronized AsyncLogWriter getWriter() {
        return writer;
    }

    public void purgeOlderThan(long txId) throws IOException {
        logManager.purgeLogsOlderThan(txId);
    }
    public CompletableFuture<Boolean> truncateJournal(DLSN dlsn) {
        return getWriter().truncate(dlsn);
    }

    public DLSN getLastDLSN() throws IOException {
        return logManager.getLastDLSN();
    }

    private void validateStoreSpec(StateStoreSpec spec) {
        checkNotNull(
            spec.getStream(),
            "No log stream is specified for state store %s", spec.getName());
    }

    private synchronized void markInitialized(AsyncLogReader reader) {
        isInitialized = true;
        // schedule periodical checkpoint
        if (null != checkpointInterval) {
            long checkpointIntervalMs = checkpointInterval.toMillis();
            checkpointTask = writeIOScheduler.scheduleAtFixedRate(
                () -> localStore.checkpoint(),
                checkpointIntervalMs,
                checkpointIntervalMs,
                TimeUnit.MILLISECONDS);
        }
        if (spec.isReadonly()) {
            replayLoop(reader);
        } else {
            reader.asyncClose();
        }
    }

    @Override
    public String name() {
        return name;
    }

    @Override
    public StateStoreSpec spec() {
        return spec;
    }

    @Override
    public CompletableFuture<Void> init(StateStoreSpec spec) {
        try {
            validateStoreSpec(spec);
        } catch (IllegalArgumentException e) {
            log.error("Fail to init state store due to : ", e);
            return FutureUtils.exception(e);
        }
        this.spec = spec;
        this.name = spec.getName();

        if (null != spec.getWriteIOScheduler()) {
            writeIOScheduler = spec.getWriteIOScheduler();
            ownWriteScheduler = false;
        } else {
            ThreadFactory threadFactory = new ThreadFactoryBuilder()
                .setNameFormat("statestore-" + spec.getName() + "-write-io-scheduler-%d")
                .build();
            writeIOScheduler = Executors.newSingleThreadScheduledExecutor(threadFactory);
            ownWriteScheduler = true;
        }

        if (null != spec.getReadIOScheduler()) {
            readIOScheduler = spec.getReadIOScheduler();
        } else if (ownWriteScheduler) {
            readIOScheduler = writeIOScheduler;
            ownReadScheduler = false;
        } else {
            ThreadFactory threadFactory = new ThreadFactoryBuilder()
                .setNameFormat("statestore-" + spec.getName() + "-read-io-scheduler-%d")
                .build();
            readIOScheduler = Executors.newSingleThreadScheduledExecutor(threadFactory);
            ownReadScheduler = true;
        }

        if (null != spec.getCheckpointStore()) {
            this.checkpointInterval = spec.getCheckpointDuration();
        } else {
            this.checkpointInterval = null;
        }

        if (spec.isReadonly()) {
            return initializeLocalStore(spec)
                .thenComposeAsync(ignored -> getLastDLSN(spec), writeIOScheduler)
                .thenComposeAsync(endDLSN -> replayJournal(endDLSN), writeIOScheduler);
        } else {
            return initializeLocalStore(spec)
                .thenComposeAsync(ignored -> initializeJournalWriter(spec), writeIOScheduler)
                .thenComposeAsync(endDLSN -> {
                    log.info("Successfully write a barrier record for mvcc store {} at {}", name, endDLSN);
                    return replayJournal(endDLSN);
                }, writeIOScheduler);
        }
    }

    private CompletableFuture<Void> initializeLocalStore(StateStoreSpec spec) {
        return executeWriteIO(() -> {
            log.info("Initializing the local state for mvcc store {}", name());
            localStore.init(spec);
            log.info("Initialized the local state for mvcc store {}", name());
            commandProcessor = newCommandProcessor();
            return null;
        });
    }

    private CompletableFuture<DLSN> initializeJournalWriter(StateStoreSpec spec) {
        synchronized (this) {
            if (null != closeFuture) {
                return FutureUtils.exception(new StateStoreClosedException(name()));
            }
        }

        try {
            logManager = logNamespace.openLog(spec.getStream());
        } catch (IOException e) {
            return FutureUtils.exception(e);
        }
        LedgerMetadata metadata = new LedgerMetadata();
        metadata.setApplication(Constants.LEDGER_METADATA_APPLICATION_STREAM_STORAGE);
        metadata.setComponent("state-store");
        return logManager.openAsyncLogWriter(metadata).thenComposeAsync(w -> {
            synchronized (this) {
                writer = w;
                nextRevision = writer.getLastTxId();
                if (nextRevision < 0) {
                    nextRevision = 0L;
                }
                log.info("Initialized the journal writer for mvcc store {} : last revision = {}",
                    name(), nextRevision);
            }
            return writeCommandBuf(newCatchupMarker());
        }, writeIOScheduler);
    }

    private CompletableFuture<DLSN> writeCatchUpMarker() {
        LedgerMetadata metadata = new LedgerMetadata();
        metadata.setApplication(Constants.LEDGER_METADATA_APPLICATION_STREAM_STORAGE);
        metadata.setComponent("state-store");
        return logManager.openAsyncLogWriter(metadata).thenComposeAsync(w -> {
            synchronized (this) {
                writer = w;
                nextRevision = writer.getLastTxId();
                if (nextRevision < 0) {
                    nextRevision = 0L;
                }
                log.info("Initialized the journal writer for writing catchup marker to mvcc store {} :"
                    + " last revision = {}", name(), nextRevision);
            }
            return writeCommandBuf(newCatchupMarker());
        }).thenCompose(dlsn -> {
            AsyncLogWriter w;
            synchronized (this) {
                w = writer;
            }
            if (null == w) {
                return FutureUtils.value(dlsn);
            } else {
                return w.asyncClose().thenApply(ignored -> dlsn);
            }
        });
    }

    private CompletableFuture<DLSN> getLastDLSN(StateStoreSpec spec) {
        synchronized (this) {
            if (null != closeFuture) {
                return FutureUtils.exception(new StateStoreClosedException(name()));
            }
        }

        try {
            logManager = logNamespace.openLog(spec.getStream());
        } catch (IOException e) {
            return FutureUtils.exception(e);
        }
        CompletableFuture<DLSN> future = FutureUtils.createFuture();
        logManager.getLastDLSNAsync().whenCompleteAsync(new FutureEventListener<DLSN>() {
            @Override
            public void onSuccess(DLSN dlsn) {
                future.complete(dlsn);
            }

            @Override
            public void onFailure(Throwable throwable) {
                if (throwable instanceof LogEmptyException || throwable instanceof LogNotFoundException) {
                    FutureUtils.proxyTo(
                        writeCatchUpMarker(),
                        future);
                } else {
                    future.completeExceptionally(throwable);
                }
            }
        });
        return future;
    }

    private CompletableFuture<Void> replayJournal(DLSN endDLSN) {
        long lastRevision = localStore.getLastRevision();
        return logManager.openAsyncLogReader(lastRevision)
            .thenComposeAsync(r -> {
                CompletableFuture<Void> replayFuture = FutureUtils.createFuture();
                replayFuture.exceptionally(
                    cause -> {
                        r.asyncClose();
                        return null;
                    });

                log.info("Successfully open the journal reader for mvcc store {} : end dlsn = {}", name(), endDLSN);
                replayJournal(r, endDLSN, replayFuture, lastRevision);
                return replayFuture;
            }, writeIOScheduler);
    }

    private void replayJournal(AsyncLogReader reader,
                               DLSN endDLSN,
                               CompletableFuture<Void> future,
                               long startingTxId) {
        synchronized (this) {
            if (null != closeFuture) {
                FutureUtils.completeExceptionally(future, new StateStoreClosedException(name()));
                return;
            }
        }

        reader.readNext().whenComplete(
            (record, exc) -> {
                if (exc != null) {
                    FutureUtils.completeExceptionally(future, exc.getCause());
                    return;
                }
                if (startingTxId != -1 && startingTxId != record.getTransactionId()) {
                    String msg = String.format(
                        "replayJournal failed: Invalid starting transaction %d expecting %d for stream %s",
                        record.getTransactionId(), startingTxId, name);
                    log.error(msg);
                    FutureUtils.completeExceptionally(future, new InvalidStateStoreException(msg));
                    return;
                }
            }
        ).whenComplete(newRecordHandler(reader, endDLSN, future));
    }

    private void replayJournal(AsyncLogReader reader,
                               DLSN endDLSN,
                               CompletableFuture<Void> future) {
        synchronized (this) {
            if (null != closeFuture) {
                FutureUtils.completeExceptionally(future, new StateStoreClosedException(name()));
                return;
            }
        }

        reader.readNext().whenComplete(newRecordHandler(reader, endDLSN, future));
    }

    private FutureEventListener<LogRecordWithDLSN> newRecordHandler(AsyncLogReader reader,
                                                                    DLSN endDLSN,
                                                                    CompletableFuture<Void> future) {
        return new FutureEventListener<LogRecordWithDLSN>() {
            @Override
            public void onSuccess(LogRecordWithDLSN record) {
                if (log.isDebugEnabled()) {
                    log.debug("Received command record {} @ {} to replay at mvcc store {}",
                        record, record.getDlsn(), name());
                }
                try {
                    if (log.isDebugEnabled()) {
                        log.debug("Applying command transaction {} - record {} @ {} to mvcc store {}",
                            record.getTransactionId(), record, record.getDlsn(), name());
                    }
                    commandProcessor.applyCommand(record.getTransactionId(), record.getPayloadBuf(), localStore);

                    if (record.getDlsn().compareTo(endDLSN) >= 0) {
                        log.info("Finished replaying journal for state store {}", name());
                        markInitialized(reader);
                        FutureUtils.complete(future, null);
                        return;
                    }

                    if (log.isDebugEnabled()) {
                        log.debug("Read next record after {} at mvcc store {}",
                            record.getDlsn(), name());
                    }
                    // read next record
                    replayJournal(reader, endDLSN, future);
                } catch (Exception e) {
                    log.error("Exception is thrown when applying command record {} @ {} to mvcc store {}",
                        record, record.getDlsn(), name());
                    FutureUtils.completeExceptionally(future, e);
                }
            }

            @Override
            public void onFailure(Throwable cause) {
                FutureUtils.completeExceptionally(future, cause);
            }
        };
    }

    private void replayLoop(AsyncLogReader reader) {
        synchronized (this) {
            if (null != closeFuture) {
                reader.asyncClose();
                return;
            }
        }

        reader.readNext().whenComplete(new FutureEventListener<LogRecordWithDLSN>() {
            @Override
            public void onSuccess(LogRecordWithDLSN record) {
                if (log.isDebugEnabled()) {
                    log.debug("Received command record {} @ {} to replay at mvcc store {}",
                        record, record.getDlsn(), name());
                }
                try {
                    commandProcessor.applyCommand(record.getTransactionId(), record.getPayloadBuf(), localStore);
                    // read next record
                    replayLoop(reader);
                } catch (StateStoreRuntimeException e) {
                    log.error("Fail to reply command record {}", record, e);
                    // TODO: handle state store exception
                }
            }

            @Override
            public void onFailure(Throwable throwable) {
                // TODO: handle read failure
            }
        });
    }

    @Override
    public CompletableFuture<Void> closeAsync() {
        CompletableFuture<Void> future;
        synchronized (this) {
            if (null != closeFuture) {
                return closeFuture;
            }
            closeFuture = future = FutureUtils.createFuture();
        }

        // cancel checkpoint task
        if (null != checkpointTask) {
            if (!checkpointTask.cancel(true)) {
                log.warn("Fail to cancel checkpoint task of state store {}", name());
            }
        }
        // wait until last checkpoint task completed
        writeIOScheduler.submit(() -> {
            log.info("closing async state store {}", name);
            FutureUtils.ensure(
                // close the log streams
                Utils.closeSequence(
                    readIOScheduler,
                    true,
                    getWriter(),
                    logManager
                ).thenRun(() -> {
                    log.info("Successfully close the log stream of state store {}", name);
                }),
                // close the local state store
                () -> {
                    if (null == readIOScheduler) {
                        closeLocalStore();
                        FutureUtils.complete(future, null);
                        return;
                    }
                    readIOScheduler.submit(() -> {
                        closeLocalStore();

                        if (ownReadScheduler) {
                            readIOScheduler.shutdown();
                        }
                        if (ownWriteScheduler) {
                            writeIOScheduler.shutdown();
                        }
                        FutureUtils.complete(future, null);
                    });
                }
            );
        });
        return future;
    }

    private void closeLocalStore() {
        if (null == localStore) {
            return;
        }
        // flush the local store
        try {
            localStore.flush();
        } catch (StateStoreException e) {
            log.warn("Fail to flush local state store {}", name(), e);
        }
        // close the local store
        localStore.close();
    }

    private <T> CompletableFuture<T> executeIO(ScheduledExecutorService scheduler,
                                               Callable<T> callable) {
        synchronized (this) {
            if (null != closeFuture) {
                return FutureUtils.exception(new StateStoreClosedException(name()));
            }
        }

        CompletableFuture<T> future = FutureUtils.createFuture();
        scheduler.submit(() -> {
            try {
                T value = callable.call();
                future.complete(value);
            } catch (Exception e) {
                future.completeExceptionally(e);
            }
        });
        return future;
    }

    //
    // Operations exposed to the implementation
    //

    protected <T> CompletableFuture<T> executeWriteIO(Callable<T> callable) {
        return executeIO(writeIOScheduler, callable);
    }

    protected <T> CompletableFuture<T> executeReadIO(Callable<T> callable) {
        return executeIO(readIOScheduler, callable);
    }

    protected synchronized CompletableFuture<DLSN> writeCommandBuf(ByteBuf cmdBuf) {
        long txId = ++nextRevision;
        return FutureUtils.ensure(
            writer.write(new LogRecord(txId, cmdBuf.nioBuffer())),
            () -> ReferenceCountUtil.release(cmdBuf));
    }

    protected synchronized CompletableFuture<Long> writeCommandBufReturnTxId(ByteBuf cmdBuf) {
        long txId = ++nextRevision;
        return FutureUtils.ensure(
            writer.write(new LogRecord(txId, cmdBuf.nioBuffer()))
                .thenApply(dlsn -> txId),
            () -> ReferenceCountUtil.release(cmdBuf));
    }

    //
    // Methods to implement
    //

    /**
     * Create a catchup marker and append the marker to the journal when initializing the journal.
     *
     * <p>The catchup marker will then be used as the end marker for replaying journal.
     *
     * <p>The implementation of this class should be able to process the catch up marker and ignore it.
     *
     * @return a catchup marker.
     */
    protected abstract ByteBuf newCatchupMarker();

    /**
     * Create a new command processor for processing the commands replayed from the journal.
     *
     * @return a new command processor.
     */
    protected abstract CommandProcessor<LocalStateStoreT> newCommandProcessor();
}
