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

import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.raft.errors.NotLeaderException;
import org.apache.kafka.server.common.OffsetAndEpoch;
import org.apache.kafka.snapshot.SnapshotReader;
import org.apache.kafka.snapshot.SnapshotWriter;

import org.slf4j.Logger;

import java.util.List;
import java.util.Optional;
import java.util.OptionalInt;

public class ReplicatedCounter implements RaftClient.Listener<Integer> {
    private static final int SNAPSHOT_DELAY_IN_RECORDS = 10;

    private final int nodeId;
    private final Logger log;
    private final RaftClient<Integer> client;

    private int committed = 0;
    private int uncommitted = 0;
    private OptionalInt claimedEpoch = OptionalInt.empty();
    private long lastOffsetSnapshotted = -1;

    private int handleLoadSnapshotCalls = 0;

    public ReplicatedCounter(
        int nodeId,
        RaftClient<Integer> client,
        LogContext logContext
    ) {
        this.nodeId = nodeId;
        this.client = client;
        log = logContext.logger(ReplicatedCounter.class);
    }

    public synchronized boolean isWritable() {
        return claimedEpoch.isPresent();
    }

    public synchronized void increment() {
        if (claimedEpoch.isEmpty()) {
            throw new KafkaException("Counter is not currently writable");
        }

        int epoch = claimedEpoch.getAsInt();
        uncommitted += 1;
        try {
            long offset = client.prepareAppend(epoch, List.of(uncommitted));
            client.schedulePreparedAppend();
            log.debug("Scheduled append of record {} with epoch {} at offset {}",
                uncommitted, epoch, offset);
        } catch (NotLeaderException e) {
            log.info("Appending failed, transition to resigned", e);
            client.resign(epoch);
        }
    }

    @Override
    public synchronized void handleCommit(BatchReader<Integer> reader) {
        try {
            int initialCommitted = committed;
            long lastCommittedOffset = -1;
            int lastCommittedEpoch = 0;
            long lastCommittedTimestamp = -1;

            while (reader.hasNext()) {
                Batch<Integer> batch = reader.next();
                log.debug(
                    "Handle commit of batch with records {} at base offset {}",
                    batch.records(),
                    batch.baseOffset()
                );
                for (Integer nextCommitted: batch.records()) {
                    if (nextCommitted != committed + 1) {
                        throw new AssertionError(
                            String.format(
                                "Expected next committed value to be %d, but instead found %d on node %d",
                                committed + 1,
                                nextCommitted,
                                nodeId
                            )
                        );
                    }
                    committed = nextCommitted;
                }

                lastCommittedOffset = batch.lastOffset();
                lastCommittedEpoch = batch.epoch();
                lastCommittedTimestamp = batch.appendTimestamp();
            }
            log.debug("Counter incremented from {} to {}", initialCommitted, committed);

            if (lastOffsetSnapshotted + SNAPSHOT_DELAY_IN_RECORDS < lastCommittedOffset) {
                log.debug(
                    "Generating new snapshot with committed offset {} and epoch {} since the previous snapshot includes {}",
                    lastCommittedOffset,
                    lastCommittedEpoch,
                    lastOffsetSnapshotted
                );
                Optional<SnapshotWriter<Integer>> snapshot = client.createSnapshot(
                    new OffsetAndEpoch(lastCommittedOffset + 1, lastCommittedEpoch),
                    lastCommittedTimestamp);
                if (snapshot.isPresent()) {
                    try {
                        snapshot.get().append(List.of(committed));
                        snapshot.get().freeze();
                        lastOffsetSnapshotted = lastCommittedOffset;
                    } finally {
                        snapshot.get().close();
                    }
                } else {
                    lastOffsetSnapshotted = lastCommittedOffset;
                }
            }
        } finally {
            reader.close();
        }
    }

    @Override
    public synchronized void handleLoadSnapshot(SnapshotReader<Integer> reader) {
        try {
            log.debug("Loading snapshot {}", reader.snapshotId());
            // Since the state machine is only one value, expect only one data record
            boolean foundDataRecord = false;
            while (reader.hasNext()) {
                Batch<Integer> batch = reader.next();
                if (!batch.records().isEmpty()) {
                    if (foundDataRecord) {
                        throw new AssertionError(
                            String.format(
                                "Expected the snapshot at %s to only one data batch %s",
                                reader.snapshotId(),
                                batch
                            )
                        );
                    } else if (batch.records().size() != 1) {
                        throw new AssertionError(
                            String.format(
                                "Expected the snapshot at %s to only contain one record %s",
                                reader.snapshotId(),
                                batch.records()
                            )
                        );
                    }

                    foundDataRecord = true;
                }

                for (Integer value : batch) {
                    log.debug("Setting value: {}", value);
                    committed = value;
                    uncommitted = value;
                }
            }
            lastOffsetSnapshotted = reader.lastContainedLogOffset();
            handleLoadSnapshotCalls += 1;
            log.debug("Finished loading snapshot. Set value: {}", committed);
        } finally {
            reader.close();
        }
    }

    @Override
    public synchronized void handleLeaderChange(LeaderAndEpoch newLeader) {
        if (newLeader.isLeader(nodeId)) {
            log.debug("Counter uncommitted value initialized to {} after claiming leadership in epoch {}",
                committed, newLeader);
            uncommitted = committed;
            claimedEpoch = OptionalInt.of(newLeader.epoch());
        } else {
            log.debug("Counter uncommitted value reset after resigning leadership");
            uncommitted = -1;
            claimedEpoch = OptionalInt.empty();
        }
        handleLoadSnapshotCalls = 0;
    }

    /** Use handleLoadSnapshotCalls to verify leader is never asked to load snapshot */
    public int handleLoadSnapshotCalls() {
        return handleLoadSnapshotCalls;
    }
}
