/*
 * 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.clients.consumer.internals;

import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.TopicAuthorizationException;
import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition;
import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic;
import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection;
import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.record.RecordBatch;
import org.apache.kafka.common.requests.AbstractRequest;
import org.apache.kafka.common.requests.OffsetsForLeaderEpochRequest;
import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;

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

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * Utility methods for preparing requests to the OffsetsForLeaderEpoch API and handling responses.
 */
public final class OffsetsForLeaderEpochUtils {

    private static final Logger LOG = LoggerFactory.getLogger(OffsetsForLeaderEpochUtils.class);

    private OffsetsForLeaderEpochUtils() {}

    static AbstractRequest.Builder<OffsetsForLeaderEpochRequest> prepareRequest(
            Map<TopicPartition, SubscriptionState.FetchPosition> requestData) {
        OffsetForLeaderTopicCollection topics = new OffsetForLeaderTopicCollection(requestData.size());
        requestData.forEach((topicPartition, fetchPosition) ->
                fetchPosition.offsetEpoch.ifPresent(fetchEpoch -> {
                    OffsetForLeaderTopic topic = topics.find(topicPartition.topic());
                    if (topic == null) {
                        topic = new OffsetForLeaderTopic().setTopic(topicPartition.topic());
                        topics.add(topic);
                    }
                    topic.partitions().add(new OffsetForLeaderPartition()
                            .setPartition(topicPartition.partition())
                            .setLeaderEpoch(fetchEpoch)
                            .setCurrentLeaderEpoch(fetchPosition.currentLeader.epoch
                                    .orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
                    );
                })
        );
        return OffsetsForLeaderEpochRequest.Builder.forConsumer(topics);
    }

    public static OffsetForEpochResult handleResponse(
            Map<TopicPartition, SubscriptionState.FetchPosition> requestData,
            OffsetsForLeaderEpochResponse response) {

        Set<TopicPartition> partitionsToRetry = new HashSet<>(requestData.keySet());
        Set<String> unauthorizedTopics = new HashSet<>();
        Map<TopicPartition, EpochEndOffset> endOffsets = new HashMap<>();

        for (OffsetForLeaderTopicResult topic : response.data().topics()) {
            for (EpochEndOffset partition : topic.partitions()) {
                TopicPartition topicPartition = new TopicPartition(topic.topic(), partition.partition());

                if (!requestData.containsKey(topicPartition)) {
                    LOG.warn("Received unrequested topic or partition {} from response, ignoring.", topicPartition);
                    continue;
                }

                Errors error = Errors.forCode(partition.errorCode());
                switch (error) {
                    case NONE:
                        LOG.debug("Handling OffsetsForLeaderEpoch response for {}. Got offset {} for epoch {}.",
                                topicPartition, partition.endOffset(), partition.leaderEpoch());
                        endOffsets.put(topicPartition, partition);
                        partitionsToRetry.remove(topicPartition);
                        break;
                    case NOT_LEADER_OR_FOLLOWER:
                    case REPLICA_NOT_AVAILABLE:
                    case KAFKA_STORAGE_ERROR:
                    case OFFSET_NOT_AVAILABLE:
                    case LEADER_NOT_AVAILABLE:
                    case FENCED_LEADER_EPOCH:
                    case UNKNOWN_LEADER_EPOCH:
                        LOG.debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
                                topicPartition, error);
                        break;
                    case UNKNOWN_TOPIC_OR_PARTITION:
                        LOG.warn("Received unknown topic or partition error in OffsetsForLeaderEpoch request for partition {}.",
                                topicPartition);
                        break;
                    case TOPIC_AUTHORIZATION_FAILED:
                        unauthorizedTopics.add(topicPartition.topic());
                        partitionsToRetry.remove(topicPartition);
                        break;
                    default:
                        LOG.warn("Attempt to fetch offsets for partition {} failed due to: {}, retrying.",
                                topicPartition, error.message());
                }
            }
        }

        if (!unauthorizedTopics.isEmpty())
            throw new TopicAuthorizationException(unauthorizedTopics);

        return new OffsetForEpochResult(endOffsets, partitionsToRetry);
    }

    static class OffsetForEpochResult {
        private final Map<TopicPartition, EpochEndOffset> endOffsets;
        private final Set<TopicPartition> partitionsToRetry;

        OffsetForEpochResult(Map<TopicPartition, EpochEndOffset> endOffsets, Set<TopicPartition> partitionsNeedingRetry) {
            this.endOffsets = endOffsets;
            this.partitionsToRetry = partitionsNeedingRetry;
        }

        public Map<TopicPartition, EpochEndOffset> endOffsets() {
            return endOffsets;
        }

        public Set<TopicPartition> partitionsToRetry() {
            return partitionsToRetry;
        }
    }
}
