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

import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.InvalidTopicException;
import org.apache.kafka.common.errors.TopicAuthorizationException;
import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.common.message.MetadataRequestData;
import org.apache.kafka.common.message.MetadataResponseData;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.requests.AbstractResponse;
import org.apache.kafka.common.requests.MetadataRequest;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.utils.LogContext;

import org.slf4j.Logger;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * Base driver implementation for APIs which target partition leaders.
 */
public class PartitionLeaderStrategy implements AdminApiLookupStrategy<TopicPartition> {
    private static final ApiRequestScope SINGLE_REQUEST_SCOPE = new ApiRequestScope() {
    };

    private final Logger log;
    private final boolean tolerateUnknownTopics;

    public PartitionLeaderStrategy(LogContext logContext) {
        this(logContext, true);
    }

    public PartitionLeaderStrategy(LogContext logContext, boolean tolerateUnknownTopics) {
        this.log = logContext.logger(PartitionLeaderStrategy.class);
        this.tolerateUnknownTopics = tolerateUnknownTopics;
    }

    @Override
    public ApiRequestScope lookupScope(TopicPartition key) {
        // Metadata requests can group topic partitions arbitrarily, so they can all share
        // the same request context
        return SINGLE_REQUEST_SCOPE;
    }

    @Override
    public MetadataRequest.Builder buildRequest(Set<TopicPartition> partitions) {
        MetadataRequestData request = new MetadataRequestData();
        request.setAllowAutoTopicCreation(false);
        partitions.stream().map(TopicPartition::topic).distinct().forEach(topic ->
            request.topics().add(new MetadataRequestData.MetadataRequestTopic().setName(topic))
        );
        return new MetadataRequest.Builder(request);
    }

    @SuppressWarnings("fallthrough")
    private void handleTopicError(
        String topic,
        Errors topicError,
        Set<TopicPartition> requestPartitions,
        Map<TopicPartition, Throwable> failed
    ) {
        switch (topicError) {
            case UNKNOWN_TOPIC_OR_PARTITION:
                if (!tolerateUnknownTopics) {
                    log.error("Received unknown topic error for topic {}", topic, topicError.exception());
                    failAllPartitionsForTopic(topic, requestPartitions, failed, tp -> topicError.exception(
                            "Failed to fetch metadata for partition " + tp + " because metadata for topic `" + topic + "` could not be found"));
                    break;
                }
            case LEADER_NOT_AVAILABLE:
            case BROKER_NOT_AVAILABLE:
                log.debug("Metadata request for topic {} returned topic-level error {}. Will retry",
                    topic, topicError);
                break;

            case TOPIC_AUTHORIZATION_FAILED:
                log.error("Received authorization failure for topic {} in `Metadata` response", topic,
                    topicError.exception());
                failAllPartitionsForTopic(topic, requestPartitions, failed, tp -> new TopicAuthorizationException(
                    "Failed to fetch metadata for partition " + tp + " due to topic authorization failure",
                    Collections.singleton(topic)));
                break;

            case INVALID_TOPIC_EXCEPTION:
                log.error("Received invalid topic error for topic {} in `Metadata` response", topic,
                    topicError.exception());
                failAllPartitionsForTopic(topic, requestPartitions, failed, tp -> new InvalidTopicException(
                    "Failed to fetch metadata for partition " + tp + " due to invalid topic `" + topic + "`",
                    Collections.singleton(topic)));
                break;

            default:
                log.error("Received unexpected error for topic {} in `Metadata` response", topic,
                    topicError.exception());
                failAllPartitionsForTopic(topic, requestPartitions, failed, tp -> topicError.exception(
                    "Failed to fetch metadata for partition " + tp + " due to unexpected error for topic `" + topic + "`"));
        }
    }

    private void failAllPartitionsForTopic(
        String topic,
        Set<TopicPartition> partitions,
        Map<TopicPartition, Throwable> failed,
        Function<TopicPartition, Throwable> exceptionGenerator
    ) {
        partitions.stream().filter(tp -> tp.topic().equals(topic)).forEach(tp ->
            failed.put(tp, exceptionGenerator.apply(tp))
        );
    }

    private void handlePartitionError(
        TopicPartition topicPartition,
        Errors partitionError,
        Map<TopicPartition, Throwable> failed
    ) {
        switch (partitionError) {
            case NOT_LEADER_OR_FOLLOWER:
            case REPLICA_NOT_AVAILABLE:
            case LEADER_NOT_AVAILABLE:
            case BROKER_NOT_AVAILABLE:
            case KAFKA_STORAGE_ERROR:
            case UNKNOWN_TOPIC_OR_PARTITION:
                log.debug("Metadata request for partition {} returned partition-level error {}. Will retry",
                    topicPartition, partitionError);
                break;

            default:
                log.error("Received unexpected error for partition {} in `Metadata` response",
                    topicPartition, partitionError.exception());
                failed.put(topicPartition, partitionError.exception(
                    "Unexpected error during metadata lookup for " + topicPartition));
        }
    }

    @Override
    public LookupResult<TopicPartition> handleResponse(
        Set<TopicPartition> requestPartitions,
        AbstractResponse abstractResponse
    ) {
        MetadataResponse response = (MetadataResponse) abstractResponse;
        Map<TopicPartition, Throwable> failed = new HashMap<>();
        Map<TopicPartition, Integer> mapped = new HashMap<>();

        for (MetadataResponseData.MetadataResponseTopic topicMetadata : response.data().topics()) {
            String topic = topicMetadata.name();
            Errors topicError = Errors.forCode(topicMetadata.errorCode());
            if (topicError != Errors.NONE) {
                handleTopicError(topic, topicError, requestPartitions, failed);
                continue;
            }

            for (MetadataResponseData.MetadataResponsePartition partitionMetadata : topicMetadata.partitions()) {
                TopicPartition topicPartition = new TopicPartition(topic, partitionMetadata.partitionIndex());
                Errors partitionError = Errors.forCode(partitionMetadata.errorCode());

                if (!requestPartitions.contains(topicPartition)) {
                    // The `Metadata` response always returns all partitions for requested
                    // topics, so we have to filter any that we are not interested in.
                    continue;
                }

                if (partitionError != Errors.NONE) {
                    handlePartitionError(topicPartition, partitionError, failed);
                    continue;
                }

                int leaderId = partitionMetadata.leaderId();
                if (leaderId >= 0) {
                    mapped.put(topicPartition, leaderId);
                } else {
                    log.debug("Metadata request for {} returned no error, but the leader is unknown. Will retry",
                        topicPartition);
                }
            }
        }
        return new LookupResult<>(failed, mapped);
    }

    /**
     * This subclass of {@link AdminApiFuture} starts with a pre-fetched map for keys to broker ids which can be
     * used to optimise the request. The map is kept up to date as metadata is fetching as this request is processed.
     * This is useful for situations in which {@link PartitionLeaderStrategy} is used
     * repeatedly, such as a sequence of identical calls to
     * {@link org.apache.kafka.clients.admin.Admin#listOffsets(Map, org.apache.kafka.clients.admin.ListOffsetsOptions)}.
     */
    public static class PartitionLeaderFuture<V> implements AdminApiFuture<TopicPartition, V> {
        private final Set<TopicPartition> requestKeys;
        private final Map<TopicPartition, Integer> partitionLeaderCache;
        private final Map<TopicPartition, KafkaFuture<V>> futures;

        public PartitionLeaderFuture(Set<TopicPartition> requestKeys, Map<TopicPartition, Integer> partitionLeaderCache) {
            this.requestKeys = requestKeys;
            this.partitionLeaderCache = partitionLeaderCache;
            this.futures = requestKeys.stream().collect(Collectors.toUnmodifiableMap(
                Function.identity(),
                k -> new KafkaFutureImpl<>()
            ));
        }

        @Override
        public Set<TopicPartition> lookupKeys() {
            return futures.keySet();
        }

        @Override
        public Set<TopicPartition> uncachedLookupKeys() {
            Set<TopicPartition> keys = new HashSet<>();
            requestKeys.forEach(tp -> {
                if (!partitionLeaderCache.containsKey(tp)) {
                    keys.add(tp);
                }
            });
            return keys;
        }

        @Override
        public Map<TopicPartition, Integer> cachedKeyBrokerIdMapping() {
            Map<TopicPartition, Integer> mapping = new HashMap<>();
            requestKeys.forEach(tp -> {
                Integer brokerId = partitionLeaderCache.get(tp);
                if (brokerId != null) {
                    mapping.put(tp, brokerId);
                }
            });
            return mapping;
        }

        public Map<TopicPartition, KafkaFuture<V>> all() {
            return futures;
        }

        @Override
        public void complete(Map<TopicPartition, V> values) {
            values.forEach(this::complete);
        }

        private void complete(TopicPartition key, V value) {
            futureOrThrow(key).complete(value);
        }

        @Override
        public void completeLookup(Map<TopicPartition, Integer> brokerIdMapping) {
            partitionLeaderCache.putAll(brokerIdMapping);
        }

        @Override
        public void completeExceptionally(Map<TopicPartition, Throwable> errors) {
            errors.forEach(this::completeExceptionally);
        }

        private void completeExceptionally(TopicPartition key, Throwable t) {
            partitionLeaderCache.remove(key);
            futureOrThrow(key).completeExceptionally(t);
        }

        private KafkaFutureImpl<V> futureOrThrow(TopicPartition key) {
            // The below typecast is safe because we initialise futures using only KafkaFutureImpl.
            KafkaFutureImpl<V> future = (KafkaFutureImpl<V>) futures.get(key);
            if (future == null) {
                throw new IllegalArgumentException("Attempt to complete future for " + key +
                    ", which was not requested");
            } else {
                return future;
            }
        }
    }
}
