/*
 * 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.kafka.image.loader;

import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.image.MetadataDelta;
import org.apache.kafka.image.MetadataImage;
import org.apache.kafka.image.MetadataProvenance;
import org.apache.kafka.image.loader.metrics.MetadataLoaderMetrics;
import org.apache.kafka.image.publisher.MetadataPublisher;
import org.apache.kafka.image.writer.ImageReWriter;
import org.apache.kafka.image.writer.ImageWriterOptions;
import org.apache.kafka.queue.EventQueue;
import org.apache.kafka.queue.KafkaEventQueue;
import org.apache.kafka.raft.Batch;
import org.apache.kafka.raft.BatchReader;
import org.apache.kafka.raft.LeaderAndEpoch;
import org.apache.kafka.raft.RaftClient;
import org.apache.kafka.server.common.ApiMessageAndVersion;
import org.apache.kafka.server.fault.FaultHandler;
import org.apache.kafka.server.fault.FaultHandlerException;
import org.apache.kafka.snapshot.SnapshotReader;
import org.apache.kafka.snapshot.Snapshots;

import org.slf4j.Logger;

import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Optional;
import java.util.OptionalLong;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

import static java.util.concurrent.TimeUnit.NANOSECONDS;


/**
 * The MetadataLoader follows changes provided by a RaftClient, and packages them into metadata
 * deltas and images that can be consumed by publishers.
 *
 * The Loader maintains its own thread, which is used to make all callbacks into publishers. If a
 * publisher A is installed before B, A will receive all callbacks before B. This is also true if
 * A and B are installed as part of a list [A, B].
 *
 * Publishers should not modify any data structures passed to them.
 *
 * It is possible to change the list of publishers dynamically over time. Whenever a new publisher is
 * added, it receives a catch-up delta which contains the full state. Any publisher installed when the
 * loader is closed will itself be closed.
 */
public class MetadataLoader implements RaftClient.Listener<ApiMessageAndVersion>, AutoCloseable {
    public static class Builder {
        private int nodeId = -1;
        private String threadNamePrefix = "";
        private Time time = Time.SYSTEM;
        private LogContext logContext = null;
        private FaultHandler faultHandler = FaultHandlerException::new;
        private MetadataLoaderMetrics metrics = null;
        private Supplier<OptionalLong> highWaterMarkAccessor = null;

        public Builder setNodeId(int nodeId) {
            this.nodeId = nodeId;
            return this;
        }

        public Builder setThreadNamePrefix(String threadNamePrefix) {
            this.threadNamePrefix = threadNamePrefix;
            return this;
        }

        public Builder setTime(Time time) {
            this.time = time;
            return this;
        }

        public Builder setFaultHandler(FaultHandler faultHandler) {
            this.faultHandler = faultHandler;
            return this;
        }

        public Builder setHighWaterMarkAccessor(Supplier<OptionalLong> highWaterMarkAccessor) {
            this.highWaterMarkAccessor = highWaterMarkAccessor;
            return this;
        }

        public Builder setMetrics(MetadataLoaderMetrics metrics) {
            this.metrics = metrics;
            return this;
        }

        public MetadataLoader build() {
            if (logContext == null) {
                logContext = new LogContext("[MetadataLoader id=" + nodeId + "] ");
            }
            if (highWaterMarkAccessor == null) {
                throw new RuntimeException("You must set the high water mark accessor.");
            }
            if (metrics == null) {
                metrics = new MetadataLoaderMetrics(Optional.empty(),
                    __ -> { },
                    __ -> { },
                    new AtomicReference<>(MetadataProvenance.EMPTY));
            }
            return new MetadataLoader(
                time,
                logContext,
                threadNamePrefix,
                faultHandler,
                metrics,
                highWaterMarkAccessor);
        }
    }

    private static final String INITIALIZE_NEW_PUBLISHERS = "InitializeNewPublishers";

    /**
     * The log4j logger for this loader.
     */
    private final Logger log;

    /**
     * The clock used by this loader.
     */
    private final Time time;

    /**
     * The fault handler to use if metadata loading fails.
     */
    private final FaultHandler faultHandler;

    /**
     * Callbacks for updating metrics.
     */
    private final MetadataLoaderMetrics metrics;

    /**
     * A function which supplies the current high water mark, or empty if it is not known.
     */
    private final Supplier<OptionalLong> highWaterMarkAccessor;

    /**
     * Publishers which haven't received any metadata yet.
     */
    private final LinkedHashMap<String, MetadataPublisher> uninitializedPublishers;

    /**
     * Publishers which are receiving updates.
     */
    private final LinkedHashMap<String, MetadataPublisher> publishers;

    /**
     * True if we have not caught up with the initial high water mark.
     * We do not send out any metadata updates until this is true.
     */
    private boolean catchingUp = true;

    /**
     * The current leader and epoch.
     */
    private LeaderAndEpoch currentLeaderAndEpoch = LeaderAndEpoch.UNKNOWN;

    /**
     * The current metadata image. Accessed only from the event queue thread.
     */
    private MetadataImage image;

    private final MetadataBatchLoader batchLoader;

    /**
     * The event queue which runs this loader.
     */
    private final KafkaEventQueue eventQueue;

    private MetadataLoader(
        Time time,
        LogContext logContext,
        String threadNamePrefix,
        FaultHandler faultHandler,
        MetadataLoaderMetrics metrics,
        Supplier<OptionalLong> highWaterMarkAccessor
    ) {
        this.log = logContext.logger(MetadataLoader.class);
        this.time = time;
        this.faultHandler = faultHandler;
        this.metrics = metrics;
        this.highWaterMarkAccessor = highWaterMarkAccessor;
        this.uninitializedPublishers = new LinkedHashMap<>();
        this.publishers = new LinkedHashMap<>();
        this.image = MetadataImage.EMPTY;
        this.batchLoader = new MetadataBatchLoader(
            logContext,
            time,
            faultHandler,
            this::maybePublishMetadata);
        this.eventQueue = new KafkaEventQueue(
            Time.SYSTEM,
            logContext,
            threadNamePrefix + "metadata-loader-",
            new ShutdownEvent());
    }

    // VisibleForTesting
    MetadataLoaderMetrics metrics() {
        return metrics;
    }

    private boolean stillNeedToCatchUp(String where, long offset) {
        if (!catchingUp) {
            log.trace("{}: we are not in the initial catching up state.", where);
            return false;
        }
        OptionalLong highWaterMark = highWaterMarkAccessor.get();
        if (highWaterMark.isEmpty()) {
            log.info("{}: the loader is still catching up because we still don't know the high " +
                    "water mark yet.", where);
            return true;
        }
        if (highWaterMark.getAsLong() - 1 > offset) {
            log.info("{}: The loader is still catching up because we have loaded up to offset " +
                    offset + ", but the high water mark is {}", where, highWaterMark.getAsLong());
            return true;
        }
        if (!batchLoader.hasSeenRecord()) {
            log.info("{}: The loader is still catching up because we have not loaded a controller record as of offset " +
                    offset + " and high water mark is {}", where, highWaterMark.getAsLong());
            return true;
        }
        log.info("{}: The loader finished catching up to the current high water mark of {}",
                where, highWaterMark.getAsLong());
        catchingUp = false;
        return false;
    }

    /**
     * Schedule an event to initialize the new publishers that are present in the system.
     *
     * @param delayNs   The minimum time in nanoseconds we should wait. If there is already an
     *                  initialization event scheduled, we will either move its deadline forward
     *                  in time or leave it unchanged.
     */
    void scheduleInitializeNewPublishers(long delayNs) {
        eventQueue.scheduleDeferred(INITIALIZE_NEW_PUBLISHERS,
            new EventQueue.EarliestDeadlineFunction(eventQueue.time().nanoseconds() + delayNs),
            () -> {
                try {
                    initializeNewPublishers();
                } catch (Throwable e) {
                    faultHandler.handleFault("Unhandled error initializing new publishers", e);
                }
            });
    }

    void initializeNewPublishers() {
        if (uninitializedPublishers.isEmpty()) {
            log.debug("InitializeNewPublishers: nothing to do.");
            return;
        }
        if (stillNeedToCatchUp("initializeNewPublishers", image.highestOffsetAndEpoch().offset())) {
            // Reschedule the initialization for later.
            log.debug("InitializeNewPublishers: unable to initialize new publisher(s) {} " +
                            "because we are still catching up with quorum metadata. Rescheduling.",
                    uninitializedPublisherNames());
            scheduleInitializeNewPublishers(TimeUnit.MILLISECONDS.toNanos(100));
            return;
        }
        log.debug("InitializeNewPublishers: setting up snapshot image for new publisher(s): {}",
                uninitializedPublisherNames());
        long startNs = time.nanoseconds();
        // We base this delta off of the empty image, reflecting the fact that these publishers
        // haven't seen anything previously.
        MetadataDelta delta = new MetadataDelta.Builder().
                setImage(MetadataImage.EMPTY).
                build();
        ImageReWriter writer = new ImageReWriter(delta);
        image.write(writer, new ImageWriterOptions.Builder(image.features().metadataVersionOrThrow()).
                setEligibleLeaderReplicasEnabled(image.features().isElrEnabled()).
                build());
        // ImageReWriter#close invokes finishSnapshot, so we don't need to invoke it here.
        SnapshotManifest manifest = new SnapshotManifest(
                image.provenance(),
                time.nanoseconds() - startNs);
        for (Iterator<MetadataPublisher> iter = uninitializedPublishers.values().iterator();
                iter.hasNext(); ) {
            MetadataPublisher publisher = iter.next();
            iter.remove();
            try {
                log.info("InitializeNewPublishers: initializing {} with a snapshot at offset {}",
                        publisher.name(), image.highestOffsetAndEpoch().offset());
                publisher.onMetadataUpdate(delta, image, manifest);
                publisher.onControllerChange(currentLeaderAndEpoch);
                publishers.put(publisher.name(), publisher);
            } catch (Throwable e) {
                faultHandler.handleFault("Unhandled error initializing " + publisher.name() +
                        " with a snapshot at offset " + image.highestOffsetAndEpoch().offset(), e);
            }
        }
    }

    private String uninitializedPublisherNames() {
        return String.join(", ", uninitializedPublishers.keySet());
    }

    /**
     * Callback used by MetadataBatchLoader and handleLoadSnapshot to update the active metadata publishers.
     */
    private void maybePublishMetadata(MetadataDelta delta, MetadataImage image, LoaderManifest manifest) {
        this.image = image;

        if (stillNeedToCatchUp(
            "maybePublishMetadata(" + manifest.type().toString() + ")",
            manifest.provenance().lastContainedOffset())
        ) {
            return;
        }

        if (log.isDebugEnabled()) {
            log.debug("handleCommit: publishing new image with provenance {}.", image.provenance());
        }
        for (MetadataPublisher publisher : publishers.values()) {
            try {
                publisher.onMetadataUpdate(delta, image, manifest);
            } catch (Throwable e) {
                faultHandler.handleFault("Unhandled error publishing the new metadata " +
                    "image ending at " + manifest.provenance().lastContainedOffset() +
                    " with publisher " + publisher.name(), e);
            }
        }
        metrics.updateLastAppliedImageProvenance(image.provenance());
        metrics.setCurrentMetadataVersion(image.features().metadataVersionOrThrow());
        if (!uninitializedPublishers.isEmpty()) {
            scheduleInitializeNewPublishers(0);
        }
    }

    @Override
    public void handleCommit(BatchReader<ApiMessageAndVersion> reader) {
        eventQueue.append(() -> {
            try (reader) {
                while (reader.hasNext()) {
                    Batch<ApiMessageAndVersion> batch = reader.next();
                    long elapsedNs = batchLoader.loadBatch(batch, currentLeaderAndEpoch);
                    metrics.updateBatchSize(batch.records().size());
                    metrics.updateBatchProcessingTimeNs(elapsedNs);
                }
                batchLoader.maybeFlushBatches(currentLeaderAndEpoch, true);
            } catch (Throwable e) {
                // This is a general catch-all block where we don't expect to end up;
                // failure-prone operations should have individual try/catch blocks around them.
                faultHandler.handleFault("Unhandled fault in MetadataLoader#handleCommit. " +
                        "Last image offset was " + image.offset(), e);
            }
        });
    }

    @Override
    public void handleLoadSnapshot(SnapshotReader<ApiMessageAndVersion> reader) {
        eventQueue.append(() -> {
            try {
                long numLoaded = metrics.incrementHandleLoadSnapshotCount();
                String snapshotName = Snapshots.filenameFromSnapshotId(reader.snapshotId());
                log.info("handleLoadSnapshot({}): incrementing HandleLoadSnapshotCount to {}.",
                    snapshotName, numLoaded);
                MetadataDelta delta = new MetadataDelta.Builder().
                    setImage(image).
                    build();
                SnapshotManifest manifest = loadSnapshot(delta, reader);
                log.info("handleLoadSnapshot({}): generated a metadata delta between offset {} " +
                        "and this snapshot in {} us.", snapshotName,
                        image.provenance().lastContainedOffset(),
                        NANOSECONDS.toMicros(manifest.elapsedNs()));
                MetadataImage image = delta.apply(manifest.provenance());
                batchLoader.resetToImage(image);
                maybePublishMetadata(delta, image, manifest);
            } catch (Throwable e) {
                // This is a general catch-all block where we don't expect to end up;
                // failure-prone operations should have individual try/catch blocks around them.
                faultHandler.handleFault("Unhandled fault in MetadataLoader#handleLoadSnapshot. " +
                        "Snapshot offset was " + reader.lastContainedLogOffset(), e);
            } finally {
                reader.close();
            }
        });
    }

    /**
     * Load a snapshot. This is relatively straightforward since we don't track as many things as
     * we do in loadLogDelta. The main complication here is that we have to maintain an index
     * of what record we are processing so that we can give useful error messages.
     *
     * @param delta     The metadata delta we are preparing.
     * @param reader    The reader which yields the snapshot batches.
     * @return          A manifest of what was loaded.
     */
    SnapshotManifest loadSnapshot(
            MetadataDelta delta,
            SnapshotReader<ApiMessageAndVersion> reader
    ) {
        long startNs = time.nanoseconds();
        int snapshotIndex = 0;
        while (reader.hasNext()) {
            Batch<ApiMessageAndVersion> batch = reader.next();
            for (ApiMessageAndVersion record : batch.records()) {
                try {
                    delta.replay(record.message());
                } catch (Throwable e) {
                    faultHandler.handleFault("Error loading metadata log record " + snapshotIndex +
                            " in snapshot at offset " + reader.lastContainedLogOffset(), e);
                }
                snapshotIndex++;
            }
        }
        delta.finishSnapshot();
        MetadataProvenance provenance = new MetadataProvenance(reader.lastContainedLogOffset(),
                reader.lastContainedLogEpoch(), reader.lastContainedLogTimestamp(), true);
        return new SnapshotManifest(provenance,
                time.nanoseconds() - startNs);
    }

    @Override
    public void handleLeaderChange(LeaderAndEpoch leaderAndEpoch) {
        eventQueue.append(() -> {
            currentLeaderAndEpoch = leaderAndEpoch;
            for (MetadataPublisher publisher : publishers.values()) {
                try {
                    publisher.onControllerChange(currentLeaderAndEpoch);
                } catch (Throwable e) {
                    faultHandler.handleFault("Unhandled error publishing the new leader " +
                        "change to " + currentLeaderAndEpoch + " with publisher " +
                        publisher.name(), e);
                }
            }
            metrics.setCurrentControllerId(leaderAndEpoch.leaderId().orElse(-1));
        });
    }

    /**
     * Install a list of publishers. When a publisher is installed, we will publish a MetadataDelta
     * to it which contains the entire current image.
     *
     * @param newPublishers     The publishers to install.
     *
     * @return                  A future which yields null when the publishers have been added, or
     *                          an exception if the installation failed.
     */
    public CompletableFuture<Void> installPublishers(List<? extends MetadataPublisher> newPublishers) {
        if (newPublishers.isEmpty()) return CompletableFuture.completedFuture(null);
        CompletableFuture<Void> future = new CompletableFuture<>();
        eventQueue.append(() -> {
            try {
                // Check that none of the publishers we are trying to install are already present.
                for (MetadataPublisher newPublisher : newPublishers) {
                    MetadataPublisher prev = publishers.get(newPublisher.name());
                    if (prev == null) {
                        prev = uninitializedPublishers.get(newPublisher.name());
                    }
                    if (prev != null) {
                        if (prev == newPublisher) {
                            throw faultHandler.handleFault("Attempted to install publisher " +
                                    newPublisher.name() + ", which is already installed.");
                        } else {
                            throw faultHandler.handleFault("Attempted to install a new publisher " +
                                    "named " + newPublisher.name() + ", but there is already a publisher " +
                                    "with that name.");
                        }
                    }
                }
                // After installation, new publishers must be initialized by sending them a full
                // snapshot of the current state. However, we cannot necessarily do that immediately,
                // because the loader itself might not be ready. Therefore, we schedule a background
                // task.
                newPublishers.forEach(p -> uninitializedPublishers.put(p.name(), p));
                scheduleInitializeNewPublishers(0);
                future.complete(null);
            } catch (Throwable e) {
                future.completeExceptionally(faultHandler.handleFault("Unhandled fault in " +
                        "MetadataLoader#installPublishers", e));
            }
        });
        return future;
    }

    // VisibleForTesting
    void waitForAllEventsToBeHandled() throws Exception {
        CompletableFuture<Void> future = new CompletableFuture<>();
        eventQueue.append(() -> future.complete(null));
        future.get();
    }

    /**
     * Remove a publisher and close it.
     *
     * @param publisher         The publisher to remove and close.
     *
     * @return                  A future which yields null when the publisher has been removed
     *                          and closed, or an exception if the removal failed.
     */
    public CompletableFuture<Void> removeAndClosePublisher(MetadataPublisher publisher) {
        CompletableFuture<Void> future = new CompletableFuture<>();
        eventQueue.append(() -> {
            try {
                if (!publishers.remove(publisher.name(), publisher)) {
                    if (!uninitializedPublishers.remove(publisher.name(), publisher)) {
                        throw faultHandler.handleFault("Attempted to remove publisher " + publisher.name() +
                                ", which is not installed.");
                    }
                }
                closePublisher(publisher);
                future.complete(null);
            } catch (Throwable e) {
                future.completeExceptionally(e);
            }
        });
        return future;
    }

    public long lastAppliedOffset() {
        return metrics.lastAppliedOffset();
    }

    @Override
    public void beginShutdown() {
        eventQueue.beginShutdown("beginShutdown");
    }

    class ShutdownEvent implements EventQueue.Event {
        @Override
        public void run() throws Exception {
            for (Iterator<MetadataPublisher> iter = uninitializedPublishers.values().iterator();
                 iter.hasNext(); ) {
                closePublisher(iter.next());
                iter.remove();
            }
            for (Iterator<MetadataPublisher> iter = publishers.values().iterator();
                 iter.hasNext(); ) {
                closePublisher(iter.next());
                iter.remove();
            }
        }
    }

    Time time() {
        return time;
    }

    private void closePublisher(MetadataPublisher publisher) {
        try {
            publisher.close();
        } catch (Throwable e) {
            faultHandler.handleFault("Got unexpected exception while closing " +
                    "publisher " + publisher.name(), e);
        }
    }

    @Override
    public void close() throws Exception {
        beginShutdown();
        eventQueue.close();
    }
}
