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

import kafka.cluster.Partition;

import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.message.FetchResponseData.PartitionData;
import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData;
import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.server.LeaderEndPoint;
import org.apache.kafka.server.PartitionFetchState;
import org.apache.kafka.server.ReplicaState;
import org.apache.kafka.server.common.CheckpointFile;
import org.apache.kafka.server.common.OffsetAndEpoch;
import org.apache.kafka.server.log.remote.storage.RemoteLogManager;
import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentMetadata;
import org.apache.kafka.server.log.remote.storage.RemoteStorageException;
import org.apache.kafka.server.log.remote.storage.RemoteStorageManager;
import org.apache.kafka.server.log.remote.storage.RemoteStorageNotReadyException;
import org.apache.kafka.storage.internals.checkpoint.LeaderEpochCheckpointFile;
import org.apache.kafka.storage.internals.log.EpochEntry;
import org.apache.kafka.storage.internals.log.LogFileUtils;
import org.apache.kafka.storage.internals.log.UnifiedLog;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static org.apache.kafka.storage.internals.log.LogStartOffsetIncrementReason.LeaderOffsetIncremented;

/**
 *  This class defines the APIs and implementation needed to handle any state transitions related to tiering
 *
 *  When started, the tier state machine will fetch the local log start offset of the
 *  leader and then build the follower's remote log aux state until the leader's
 *  local log start offset.
 */
public class TierStateMachine {
    private static final Logger log = LoggerFactory.getLogger(TierStateMachine.class);

    private final LeaderEndPoint leader;
    private final ReplicaManager replicaMgr;
    private final boolean useFutureLog;
    public TierStateMachine(LeaderEndPoint leader,
                            ReplicaManager replicaMgr,
                            boolean useFutureLog) {
        this.leader = leader;
        this.replicaMgr = replicaMgr;
        this.useFutureLog = useFutureLog;
    }

    /**
     * Start the tier state machine for the provided topic partition.
     *
     * @param topicPartition the topic partition
     * @param currentFetchState the current PartitionFetchState which will
     *                          be used to derive the return value
     * @param fetchPartitionData the data from the fetch response that returned the offset moved to tiered storage error
     *
     * @return the new PartitionFetchState after the successful start of the
     *         tier state machine
     */
    PartitionFetchState start(TopicPartition topicPartition,
                              PartitionFetchState currentFetchState,
                              PartitionData fetchPartitionData) throws Exception {
        OffsetAndEpoch epochAndLeaderLocalStartOffset = leader.fetchEarliestLocalOffset(topicPartition, currentFetchState.currentLeaderEpoch());
        int epoch = epochAndLeaderLocalStartOffset.epoch();
        long leaderLocalStartOffset = epochAndLeaderLocalStartOffset.offset();

        long offsetToFetch;
        replicaMgr.brokerTopicStats().topicStats(topicPartition.topic()).buildRemoteLogAuxStateRequestRate().mark();
        replicaMgr.brokerTopicStats().allTopicsStats().buildRemoteLogAuxStateRequestRate().mark();

        UnifiedLog unifiedLog;
        if (useFutureLog) {
            unifiedLog = replicaMgr.futureLogOrException(topicPartition);
        } else {
            unifiedLog = replicaMgr.localLogOrException(topicPartition);
        }

        try {
            offsetToFetch = buildRemoteLogAuxState(topicPartition, currentFetchState.currentLeaderEpoch(), leaderLocalStartOffset, epoch, fetchPartitionData.logStartOffset(), unifiedLog);
        } catch (RemoteStorageException e) {
            replicaMgr.brokerTopicStats().topicStats(topicPartition.topic()).failedBuildRemoteLogAuxStateRate().mark();
            replicaMgr.brokerTopicStats().allTopicsStats().failedBuildRemoteLogAuxStateRate().mark();
            throw e;
        }

        OffsetAndEpoch fetchLatestOffsetResult = leader.fetchLatestOffset(topicPartition, currentFetchState.currentLeaderEpoch());
        long leaderEndOffset = fetchLatestOffsetResult.offset();

        long initialLag = leaderEndOffset - offsetToFetch;

        return new PartitionFetchState(currentFetchState.topicId(), offsetToFetch, Optional.of(initialLag), currentFetchState.currentLeaderEpoch(),
                ReplicaState.FETCHING, unifiedLog.latestEpoch());

    }

    private OffsetForLeaderEpochResponseData.EpochEndOffset fetchEarlierEpochEndOffset(Integer epoch,
                                                                                       TopicPartition partition,
                                                                                       Integer currentLeaderEpoch) {
        int previousEpoch = epoch - 1;

        // Find the end-offset for the epoch earlier to the given epoch from the leader
        Map<TopicPartition, OffsetForLeaderEpochRequestData.OffsetForLeaderPartition> partitionsWithEpochs = new HashMap<>();
        partitionsWithEpochs.put(partition, new OffsetForLeaderEpochRequestData.OffsetForLeaderPartition().setPartition(partition.partition()).setCurrentLeaderEpoch(currentLeaderEpoch).setLeaderEpoch(previousEpoch));
        var epochEndOffset = leader.fetchEpochEndOffsets(partitionsWithEpochs).get(partition);

        if (epochEndOffset == null) {
            throw new KafkaException("No response received for partition: " + partition);
        }

        if (epochEndOffset.errorCode() != Errors.NONE.code()) {
            throw Errors.forCode(epochEndOffset.errorCode()).exception();
        }

        return epochEndOffset;
    }

    private List<EpochEntry> readLeaderEpochCheckpoint(RemoteLogManager rlm,
                                                       RemoteLogSegmentMetadata remoteLogSegmentMetadata) throws IOException, RemoteStorageException {
        try (InputStream inputStream = rlm.storageManager().fetchIndex(remoteLogSegmentMetadata, RemoteStorageManager.IndexType.LEADER_EPOCH);
             BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            CheckpointFile.CheckpointReadBuffer<EpochEntry> readBuffer = new CheckpointFile.CheckpointReadBuffer<>("", bufferedReader, 0, LeaderEpochCheckpointFile.FORMATTER);
            return readBuffer.read();
        }
    }

    private void buildProducerSnapshotFile(UnifiedLog unifiedLog,
                                           long nextOffset,
                                           RemoteLogSegmentMetadata remoteLogSegmentMetadata,
                                           RemoteLogManager rlm) throws IOException, RemoteStorageException {
        // Restore producer snapshot
        File snapshotFile = LogFileUtils.producerSnapshotFile(unifiedLog.dir(), nextOffset);
        Path tmpSnapshotFile = Paths.get(snapshotFile.getAbsolutePath() + ".tmp");
        // Copy it to snapshot file in atomic manner.
        try (InputStream inputStream = rlm.storageManager().fetchIndex(remoteLogSegmentMetadata, RemoteStorageManager.IndexType.PRODUCER_SNAPSHOT)) {
            Files.copy(inputStream, tmpSnapshotFile, StandardCopyOption.REPLACE_EXISTING);
        }
        Utils.atomicMoveWithFallback(tmpSnapshotFile, snapshotFile.toPath(), false);

        // Reload producer snapshots.
        unifiedLog.producerStateManager().truncateFullyAndReloadSnapshots();
        unifiedLog.loadProducerState(nextOffset);
    }

    /**
     * It tries to build the required state for this partition from leader and remote storage so that it can start
     * fetching records from the leader. The return value is the next offset to fetch from the leader, which is the
     * next offset following the end offset of the remote log portion.
     */
    private Long buildRemoteLogAuxState(TopicPartition topicPartition,
                                        Integer currentLeaderEpoch,
                                        Long leaderLocalLogStartOffset,
                                        Integer epochForLeaderLocalLogStartOffset,
                                        Long leaderLogStartOffset,
                                        UnifiedLog unifiedLog) throws IOException, RemoteStorageException {

        if (!unifiedLog.remoteLogEnabled()) {
            // If the tiered storage is not enabled throw an exception back so that it will retry until the tiered storage
            // is set as expected.
            throw new RemoteStorageException("Couldn't build the state from remote store for partition " + topicPartition + ", as remote log storage is not yet enabled");
        }

        if (replicaMgr.remoteLogManager().isEmpty())
            throw new IllegalStateException("RemoteLogManager is not yet instantiated");

        RemoteLogManager rlm = replicaMgr.remoteLogManager().get();

        // Find the respective leader epoch for (leaderLocalLogStartOffset - 1). We need to build the leader epoch cache
        // until that offset
        long previousOffsetToLeaderLocalLogStartOffset = leaderLocalLogStartOffset - 1;
        int targetEpoch;
        // If the existing epoch is 0, no need to fetch from earlier epoch as the desired offset(leaderLogStartOffset - 1)
        // will have the same epoch.
        if (epochForLeaderLocalLogStartOffset == 0) {
            targetEpoch = epochForLeaderLocalLogStartOffset;
        } else {
            // Fetch the earlier epoch/end-offset(exclusive) from the leader.
            OffsetForLeaderEpochResponseData.EpochEndOffset earlierEpochEndOffset = fetchEarlierEpochEndOffset(epochForLeaderLocalLogStartOffset, topicPartition, currentLeaderEpoch);
            // Check if the target offset lies within the range of earlier epoch. Here, epoch's end-offset is exclusive.
            if (earlierEpochEndOffset.endOffset() > previousOffsetToLeaderLocalLogStartOffset) {
                // Always use the leader epoch from returned earlierEpochEndOffset.
                // This gives the respective leader epoch, that will handle any gaps in epochs.
                // For ex, leader epoch cache contains:
                // leader-epoch   start-offset
                //  0               20
                //  1               85
                //  <2> - gap no messages were appended in this leader epoch.
                //  3               90
                //  4               98
                // There is a gap in leader epoch. For leaderLocalLogStartOffset as 90, leader-epoch is 3.
                // fetchEarlierEpochEndOffset(2) will return leader-epoch as 1, end-offset as 90.
                // So, for offset 89, we should return leader epoch as 1 like below.
                targetEpoch = earlierEpochEndOffset.leaderEpoch();
            } else {
                targetEpoch = epochForLeaderLocalLogStartOffset;
            }
        }

        if (!rlm.isPartitionReady(topicPartition)) {
            throw new RemoteStorageNotReadyException("RemoteLogManager is not ready for partition: " + topicPartition);
        }

        RemoteLogSegmentMetadata remoteLogSegmentMetadata = rlm.fetchRemoteLogSegmentMetadata(topicPartition, targetEpoch, previousOffsetToLeaderLocalLogStartOffset)
                .orElseThrow(() -> buildRemoteStorageException(topicPartition, targetEpoch, currentLeaderEpoch,
                        leaderLocalLogStartOffset, leaderLogStartOffset));


        // Build leader epoch cache, producer snapshots until remoteLogSegmentMetadata.endOffset() and start
        // segments from (remoteLogSegmentMetadata.endOffset() + 1)
        // Assign nextOffset with the offset from which next fetch should happen.
        long nextOffset = remoteLogSegmentMetadata.endOffset() + 1;

        // Truncate the existing local log before restoring the leader epoch cache and producer snapshots.
        Partition partition = replicaMgr.getPartitionOrException(topicPartition);
        partition.truncateFullyAndStartAt(nextOffset, useFutureLog, Optional.of(leaderLogStartOffset));
        // Increment start offsets
        unifiedLog.maybeIncrementLogStartOffset(leaderLogStartOffset, LeaderOffsetIncremented);
        unifiedLog.maybeIncrementLocalLogStartOffset(nextOffset, LeaderOffsetIncremented);

        // Build leader epoch cache.
        List<EpochEntry> epochs = readLeaderEpochCheckpoint(rlm, remoteLogSegmentMetadata);
        unifiedLog.leaderEpochCache().assign(epochs);

        log.info("Updated the epoch cache from remote tier till offset: {} with size: {} for {}", leaderLocalLogStartOffset, epochs.size(), partition);

        buildProducerSnapshotFile(unifiedLog, nextOffset, remoteLogSegmentMetadata, rlm);

        log.debug("Built the leader epoch cache and producer snapshots from remote tier for {}, " +
                        "with active producers size: {}, leaderLogStartOffset: {}, and logEndOffset: {}",
                partition, unifiedLog.producerStateManager().activeProducers().size(), leaderLogStartOffset, nextOffset);

        return nextOffset;
    }

    private RemoteStorageException buildRemoteStorageException(TopicPartition topicPartition,
                                                               int targetEpoch,
                                                               int currentLeaderEpoch,
                                                               long leaderLocalLogStartOffset,
                                                               long leaderLogStartOffset) {
        String message = String.format(
                "Couldn't build the state from remote store for partition: %s, currentLeaderEpoch: %d, " +
                        "leaderLocalLogStartOffset: %d, leaderLogStartOffset: %d, epoch: %d as the previous remote log segment metadata was not found",
                topicPartition, currentLeaderEpoch, leaderLocalLogStartOffset, leaderLogStartOffset, targetEpoch
        );
        return new RemoteStorageException(message);
    }
}
