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

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.loader.LoaderManifest;
import org.apache.kafka.image.loader.LogDeltaManifest;
import org.apache.kafka.queue.EventQueue;
import org.apache.kafka.queue.KafkaEventQueue;
import org.apache.kafka.server.fault.FaultHandler;

import org.slf4j.Logger;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;


/**
 * A metadata publisher that generates snapshots when appropriate.
 */
public class SnapshotGenerator implements MetadataPublisher {
    public static class Builder {
        private final Emitter emitter;
        private int nodeId = 0;
        private Time time = Time.SYSTEM;
        private FaultHandler faultHandler = (m, e) -> null;
        private long maxBytesSinceLastSnapshot = 100 * 1024L * 1024L;
        private long maxTimeSinceLastSnapshotNs = TimeUnit.DAYS.toNanos(1);
        private AtomicReference<String> disabledReason = null;
        private String threadNamePrefix = "";

        public Builder(Emitter emitter) {
            this.emitter = emitter;
        }

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

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

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

        public Builder setMaxBytesSinceLastSnapshot(long maxBytesSinceLastSnapshot) {
            this.maxBytesSinceLastSnapshot = maxBytesSinceLastSnapshot;
            return this;
        }

        public Builder setMaxTimeSinceLastSnapshotNs(long maxTimeSinceLastSnapshotNs) {
            this.maxTimeSinceLastSnapshotNs = maxTimeSinceLastSnapshotNs;
            return this;
        }

        public Builder setDisabledReason(AtomicReference<String> disabledReason) {
            this.disabledReason = disabledReason;
            return this;
        }

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

        public SnapshotGenerator build() {
            if (disabledReason == null) {
                disabledReason = new AtomicReference<>();
            }
            return new SnapshotGenerator(
                nodeId,
                time,
                emitter,
                faultHandler,
                maxBytesSinceLastSnapshot,
                maxTimeSinceLastSnapshotNs,
                disabledReason,
                threadNamePrefix
            );
        }
    }

    /**
     * The callback which actually generates the snapshot.
     */
    public interface Emitter {
        /**
         * Emit a snapshot for the given image.
         *
         * Note: if a snapshot has already been emitted for the given offset and epoch pair, this
         * function will not recreate it.
         *
         * @param image     The metadata image to emit.
         */
        void maybeEmit(MetadataImage image);
    }

    /**
     * The node ID.
     */
    private final int nodeId;

    /**
     * The clock to use.
     */
    private final Time time;

    /**
     * The emitter callback, which actually generates the snapshot.
     */
    private final Emitter emitter;

    /**
     * The slf4j logger to use.
     */
    private final Logger log;

    /**
     * The fault handler to use.
     */
    private final FaultHandler faultHandler;

    /**
     * The maximum number of bytes we will wait to see before triggering a new snapshot.
     */
    private final long maxBytesSinceLastSnapshot;

    /**
     * The maximum amount of time we will wait before triggering a snapshot, or 0 to disable
     * time-based snapshotting.
     */
    private final long maxTimeSinceLastSnapshotNs;

    /**
     * If non-null, the reason why snapshots have been disabled.
     */
    private final AtomicReference<String> disabledReason;

    /**
     * The event queue used to schedule emitting snapshots.
     */
    private final EventQueue eventQueue;

    /**
     * The log bytes that we have read since the last snapshot.
     */
    private long bytesSinceLastSnapshot;

    /**
     * The time at which we created the last snapshot.
     */
    private long lastSnapshotTimeNs;

    private SnapshotGenerator(
        int nodeId,
        Time time,
        Emitter emitter,
        FaultHandler faultHandler,
        long maxBytesSinceLastSnapshot,
        long maxTimeSinceLastSnapshotNs,
        AtomicReference<String> disabledReason,
        String threadNamePrefix
    ) {
        this.nodeId = nodeId;
        this.time = time;
        this.emitter = emitter;
        this.faultHandler = faultHandler;
        this.maxBytesSinceLastSnapshot = maxBytesSinceLastSnapshot;
        this.maxTimeSinceLastSnapshotNs = maxTimeSinceLastSnapshotNs;
        LogContext logContext = new LogContext("[SnapshotGenerator id=" + nodeId + "] ");
        this.log = logContext.logger(SnapshotGenerator.class);
        this.disabledReason = disabledReason;
        this.eventQueue = new KafkaEventQueue(time, logContext, threadNamePrefix + "snapshot-generator-");
        resetSnapshotCounters();
        log.debug("Starting SnapshotGenerator.");
    }

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

    void resetSnapshotCounters() {
        this.bytesSinceLastSnapshot = 0L;
        this.lastSnapshotTimeNs = time.nanoseconds();
    }

    @Override
    public void onMetadataUpdate(
        MetadataDelta delta,
        MetadataImage newImage,
        LoaderManifest manifest
    ) {
        switch (manifest.type()) {
            case LOG_DELTA:
                publishLogDelta(newImage, (LogDeltaManifest) manifest);
                break;
            case SNAPSHOT:
                publishSnapshot(newImage);
                break;
        }
    }

    void publishSnapshot(MetadataImage newImage) {
        log.debug("Resetting the snapshot counters because we just read {}.", newImage.provenance().snapshotName());
        resetSnapshotCounters();
    }

    void publishLogDelta(MetadataImage newImage, LogDeltaManifest manifest) {
        bytesSinceLastSnapshot += manifest.numBytes();
        if (bytesSinceLastSnapshot >= maxBytesSinceLastSnapshot) {
            if (eventQueue.isEmpty()) {
                maybeScheduleEmit("we have replayed at least " + maxBytesSinceLastSnapshot +
                    " bytes", newImage, manifest.provenance().isOffsetBatchAligned());
            } else if (log.isTraceEnabled()) {
                log.trace("Not scheduling bytes-based snapshot because event queue is not empty yet.");
            }
        } else if (maxTimeSinceLastSnapshotNs != 0 &&
                (time.nanoseconds() - lastSnapshotTimeNs >= maxTimeSinceLastSnapshotNs)) {
            if (eventQueue.isEmpty()) {
                maybeScheduleEmit("we have waited at least " +
                    TimeUnit.NANOSECONDS.toMinutes(maxTimeSinceLastSnapshotNs) +
                    " minute(s)", newImage, manifest.provenance().isOffsetBatchAligned());
            } else if (log.isTraceEnabled()) {
                log.trace("Not scheduling time-based snapshot because event queue is not empty yet.");
            }
        } else if (log.isTraceEnabled()) {
            log.trace("Neither time-based nor bytes-based criteria are met; not scheduling snapshot.");
        }
    }

    void maybeScheduleEmit(
        String reason,
        MetadataImage image,
        boolean isOffsetBatchAligned
    ) {
        String currentDisabledReason = disabledReason.get();
        if (currentDisabledReason != null) {
            log.error("Not emitting {} despite the fact that {} because snapshots are " +
                "disabled; {}", image.provenance().snapshotName(), reason, currentDisabledReason);
        } else if (!isOffsetBatchAligned) {
            log.debug("Not emitting {} despite the fact that {} because snapshots are " +
                "disabled; {}", image.provenance().snapshotName(), reason, "metadata image is not batch aligned");
        } else {
            eventQueue.append(() -> {
                resetSnapshotCounters();
                log.info("Creating new KRaft snapshot file {} because {}.",
                        image.provenance().snapshotName(), reason);
                try {
                    emitter.maybeEmit(image);
                } catch (Throwable e) {
                    faultHandler.handleFault("KRaft snapshot file generation error", e);
                }
            });
        }
    }

    public void beginShutdown() {
        log.debug("Beginning shutdown of SnapshotGenerator.");
        this.disabledReason.compareAndSet(null, "we are shutting down");
        eventQueue.beginShutdown("beginShutdown");
    }

    @Override
    public void close() throws InterruptedException {
        eventQueue.beginShutdown("close");
        log.debug("Closing SnapshotGenerator.");
        eventQueue.close();
    }
}
