/*
 * 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.pulsar.metadata.impl.batching;

import com.google.common.util.concurrent.MoreExecutors;
import io.netty.util.concurrent.DefaultThreadFactory;
import java.util.Collections;
import java.util.EnumSet;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.metadata.api.GetResult;
import org.apache.pulsar.metadata.api.MetadataEventSynchronizer;
import org.apache.pulsar.metadata.api.MetadataStoreConfig;
import org.apache.pulsar.metadata.api.MetadataStoreException;
import org.apache.pulsar.metadata.api.Stat;
import org.apache.pulsar.metadata.api.extended.CreateOption;
import org.apache.pulsar.metadata.impl.AbstractMetadataStore;
import org.apache.pulsar.metadata.impl.stats.BatchMetadataStoreStats;
import org.jctools.queues.MessagePassingQueue;
import org.jctools.queues.MpscUnboundedArrayQueue;
import org.jspecify.annotations.Nullable;

@Slf4j
public abstract class AbstractBatchedMetadataStore extends AbstractMetadataStore {

    private final ScheduledFuture<?> scheduledTask;
    private final MessagePassingQueue<MetadataOp> readOps;
    private final MessagePassingQueue<MetadataOp> writeOps;

    private final boolean enabled;
    private final int maxDelayMillis;
    protected final int maxOperations;
    protected int maxSize;
    private MetadataEventSynchronizer synchronizer;
    private final BatchMetadataStoreStats batchMetadataStoreStats;
    protected MetadataStoreBatchStrategy metadataStoreBatchStrategy;
    @Nullable
    private final ScheduledExecutorService flushExecutor;

    protected AbstractBatchedMetadataStore(MetadataStoreConfig conf) {
        super(conf.getMetadataStoreName(), conf.getOpenTelemetry(), conf.getNodeSizeStats(),
                conf.getNumSerDesThreads());

        this.enabled = conf.isBatchingEnabled();
        this.maxDelayMillis = conf.getBatchingMaxDelayMillis();
        this.maxOperations = conf.getBatchingMaxOperations();
        this.maxSize = conf.getBatchingMaxSizeKb() * 1_024;

        if (enabled) {
            readOps = new MpscUnboundedArrayQueue<>(10_000);
            writeOps = new MpscUnboundedArrayQueue<>(10_000);
            final var name = StringUtils.isNotBlank(conf.getMetadataStoreName()) ? conf.getMetadataStoreName()
                    : getClass().getSimpleName();
            flushExecutor = Executors.newSingleThreadScheduledExecutor(new DefaultThreadFactory(
                    name + "-batch-flusher"));
            scheduledTask = flushExecutor.scheduleAtFixedRate(this::flush, maxDelayMillis, maxDelayMillis,
                    TimeUnit.MILLISECONDS);
        } else {
            scheduledTask = null;
            readOps = null;
            writeOps = null;
            flushExecutor = null;
        }

        // update synchronizer and register sync listener
        updateMetadataEventSynchronizer(conf.getSynchronizer());
        this.batchMetadataStoreStats = new BatchMetadataStoreStats(metadataStoreName);
        this.metadataStoreBatchStrategy = new DefaultMetadataStoreBatchStrategy(maxOperations, maxSize);
    }

    @Override
    public void close() throws Exception {
        if (enabled) {
            // Fail all the pending items
            MetadataStoreException ex =
                    new MetadataStoreException.AlreadyClosedException("Metadata store is getting closed");
            MetadataOp op;
            while ((op = readOps.poll()) != null) {
                op.getFuture().completeExceptionally(ex);
            }
            while ((op = writeOps.poll()) != null) {
                op.getFuture().completeExceptionally(ex);
            }
            scheduledTask.cancel(true);
            MoreExecutors.shutdownAndAwaitTermination(flushExecutor, 10, TimeUnit.SECONDS);
        }
        super.close();
        this.batchMetadataStoreStats.close();
    }

    private synchronized void flush() {
        List<MetadataOp> currentBatch;
        if (!readOps.isEmpty()) {
            while (CollectionUtils.isNotEmpty(currentBatch = metadataStoreBatchStrategy.nextBatch(readOps))) {
                internalBatchOperation(currentBatch);
            }
        }
        if (!writeOps.isEmpty()) {
            while (CollectionUtils.isNotEmpty(currentBatch = metadataStoreBatchStrategy.nextBatch(writeOps))) {
                internalBatchOperation(currentBatch);
            }
        }
    }

    @Override
    public final CompletableFuture<Optional<GetResult>> storeGet(String path) {
        OpGet op = new OpGet(path);
        enqueue(readOps, op);
        return op.getFuture();
    }

    @Override
    public final CompletableFuture<List<String>> getChildrenFromStore(String path) {
        OpGetChildren op = new OpGetChildren(path);
        enqueue(readOps, op);
        return op.getFuture();
    }

    @Override
    protected final CompletableFuture<Void> storeDelete(String path, Optional<Long> expectedVersion) {
        OpDelete op = new OpDelete(path, expectedVersion);
        enqueue(writeOps, op);
        return op.getFuture();
    }

    @Override
    protected CompletableFuture<Stat> storePut(String path, byte[] data, Optional<Long> optExpectedVersion,
                                               EnumSet<CreateOption> options) {
        OpPut op = new OpPut(path, data, optExpectedVersion, options);
        enqueue(writeOps, op);
        return op.getFuture();
    }

    @Override
    public Optional<MetadataEventSynchronizer> getMetadataEventSynchronizer() {
        return Optional.ofNullable(synchronizer);
    }

    @Override
    public void updateMetadataEventSynchronizer(MetadataEventSynchronizer synchronizer) {
        this.synchronizer = synchronizer;
        registerSyncListener(Optional.ofNullable(synchronizer));
    }

    private void enqueue(MessagePassingQueue<MetadataOp> queue, MetadataOp op) {
        if (isClosed()) {
            MetadataStoreException ex = new MetadataStoreException.AlreadyClosedException();
            op.getFuture().completeExceptionally(ex);
            return;
        }
        if (enabled) {
            if (!queue.offer(op)) {
                // Execute individually if we're failing to enqueue
                internalBatchOperation(Collections.singletonList(op));
                return;
            }
            if (queue.size() > maxOperations) {
                flush();
            }
        } else {
            internalBatchOperation(Collections.singletonList(op));
        }
    }

    private void internalBatchOperation(List<MetadataOp> ops) {
        if (isClosed()) {
            MetadataStoreException ex =
                    new MetadataStoreException.AlreadyClosedException();
            ops.forEach(op -> op.getFuture().completeExceptionally(ex));
            return;
        }
        long now = System.currentTimeMillis();
        for (MetadataOp op : ops) {
            this.batchMetadataStoreStats.recordOpWaiting(now - op.created());
        }
        this.batchOperation(ops);
        this.batchMetadataStoreStats.recordOpsInBatch(ops.size());
        this.batchMetadataStoreStats.recordBatchExecuteTime(System.currentTimeMillis() - now);
    }

    protected abstract void batchOperation(List<MetadataOp> ops);

    protected final void safeExecuteCallbacks(Runnable runnable, List<MetadataOp> ops) {
        safeExecuteCallback(runnable, t -> ops.forEach(op -> op.getFuture().completeExceptionally(t)));
    }
}
