/*
 * 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.clients.admin.DescribeProducersOptions;
import org.apache.kafka.clients.admin.DescribeProducersResult.PartitionProducerState;
import org.apache.kafka.clients.admin.ProducerState;
import org.apache.kafka.common.Node;
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.message.DescribeProducersRequestData;
import org.apache.kafka.common.message.DescribeProducersResponseData;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.requests.AbstractResponse;
import org.apache.kafka.common.requests.ApiError;
import org.apache.kafka.common.requests.DescribeProducersRequest;
import org.apache.kafka.common.requests.DescribeProducersResponse;
import org.apache.kafka.common.utils.CollectionUtils;
import org.apache.kafka.common.utils.LogContext;

import org.slf4j.Logger;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.OptionalInt;
import java.util.OptionalLong;
import java.util.Set;
import java.util.stream.Collectors;

public class DescribeProducersHandler extends AdminApiHandler.Batched<TopicPartition, PartitionProducerState> {
    private final Logger log;
    private final DescribeProducersOptions options;
    private final AdminApiLookupStrategy<TopicPartition> lookupStrategy;

    public DescribeProducersHandler(
        DescribeProducersOptions options,
        LogContext logContext
    ) {
        this.options = options;
        this.log = logContext.logger(DescribeProducersHandler.class);

        if (options.brokerId().isPresent()) {
            this.lookupStrategy = new StaticBrokerStrategy<>(options.brokerId().getAsInt());
        } else {
            this.lookupStrategy = new PartitionLeaderStrategy(logContext);
        }
    }

    public static PartitionLeaderStrategy.PartitionLeaderFuture<PartitionProducerState> newFuture(
        Collection<TopicPartition> topicPartitions,
        Map<TopicPartition, Integer> partitionLeaderCache
    ) {
        return new PartitionLeaderStrategy.PartitionLeaderFuture<>(new HashSet<>(topicPartitions), partitionLeaderCache);
    }

    @Override
    public String apiName() {
        return "describeProducers";
    }

    @Override
    public AdminApiLookupStrategy<TopicPartition> lookupStrategy() {
        return lookupStrategy;
    }

    @Override
    public DescribeProducersRequest.Builder buildBatchedRequest(
        int brokerId,
        Set<TopicPartition> topicPartitions
    ) {
        DescribeProducersRequestData request = new DescribeProducersRequestData();
        DescribeProducersRequest.Builder builder = new DescribeProducersRequest.Builder(request);

        CollectionUtils.groupPartitionsByTopic(
            topicPartitions,
            builder::addTopic,
            (topicRequest, partitionId) -> topicRequest.partitionIndexes().add(partitionId)
        );

        return builder;
    }

    private void handlePartitionError(
        TopicPartition topicPartition,
        ApiError apiError,
        Map<TopicPartition, Throwable> failed,
        List<TopicPartition> unmapped
    ) {
        switch (apiError.error()) {
            case NOT_LEADER_OR_FOLLOWER:
                if (options.brokerId().isPresent()) {
                    // Typically these errors are retriable, but if the user specified the brokerId
                    // explicitly, then they are fatal.
                    int brokerId = options.brokerId().getAsInt();
                    log.error("Not leader error in `DescribeProducers` response for partition {} " +
                        "for brokerId {} set in options", topicPartition, brokerId, apiError.exception());
                    failed.put(topicPartition, apiError.error().exception("Failed to describe active producers " +
                        "for partition " + topicPartition + " on brokerId " + brokerId));
                } else {
                    // Otherwise, we unmap the partition so that we can find the new leader
                    log.debug("Not leader error in `DescribeProducers` response for partition {}. " +
                        "Will retry later.", topicPartition);
                    unmapped.add(topicPartition);
                }
                break;

            case UNKNOWN_TOPIC_OR_PARTITION:
                log.debug("Unknown topic/partition error in `DescribeProducers` response for partition {}. " +
                    "Will retry later.", topicPartition);
                break;

            case INVALID_TOPIC_EXCEPTION:
                log.error("Invalid topic in `DescribeProducers` response for partition {}",
                    topicPartition, apiError.exception());
                failed.put(topicPartition, new InvalidTopicException(
                    "Failed to fetch metadata for partition " + topicPartition
                        + " due to invalid topic error: " + apiError.messageWithFallback(),
                    Collections.singleton(topicPartition.topic())));
                break;

            case TOPIC_AUTHORIZATION_FAILED:
                log.error("Authorization failed in `DescribeProducers` response for partition {}",
                    topicPartition, apiError.exception());
                failed.put(topicPartition, new TopicAuthorizationException("Failed to describe " +
                    "active producers for partition " + topicPartition + " due to authorization failure on topic" +
                    " `" + topicPartition.topic() + "`", Collections.singleton(topicPartition.topic())));
                break;

            default:
                log.error("Unexpected error in `DescribeProducers` response for partition {}",
                    topicPartition, apiError.exception());
                failed.put(topicPartition, apiError.error().exception("Failed to describe active " +
                    "producers for partition " + topicPartition + " due to unexpected error"));
                break;
        }
    }

    @Override
    public ApiResult<TopicPartition, PartitionProducerState> handleResponse(
        Node broker,
        Set<TopicPartition> keys,
        AbstractResponse abstractResponse
    ) {
        DescribeProducersResponse response = (DescribeProducersResponse) abstractResponse;
        Map<TopicPartition, PartitionProducerState> completed = new HashMap<>();
        Map<TopicPartition, Throwable> failed = new HashMap<>();
        List<TopicPartition> unmapped = new ArrayList<>();

        for (DescribeProducersResponseData.TopicResponse topicResponse : response.data().topics()) {
            for (DescribeProducersResponseData.PartitionResponse partitionResponse : topicResponse.partitions()) {
                TopicPartition topicPartition = new TopicPartition(
                    topicResponse.name(), partitionResponse.partitionIndex());

                Errors error = Errors.forCode(partitionResponse.errorCode());
                if (error != Errors.NONE) {
                    ApiError apiError = new ApiError(error, partitionResponse.errorMessage());
                    handlePartitionError(topicPartition, apiError, failed, unmapped);
                    continue;
                }

                List<ProducerState> activeProducers = partitionResponse.activeProducers().stream()
                    .map(activeProducer -> {
                        OptionalLong currentTransactionFirstOffset =
                            activeProducer.currentTxnStartOffset() < 0 ?
                                OptionalLong.empty() :
                                OptionalLong.of(activeProducer.currentTxnStartOffset());
                        OptionalInt coordinatorEpoch =
                            activeProducer.coordinatorEpoch() < 0 ?
                                OptionalInt.empty() :
                                OptionalInt.of(activeProducer.coordinatorEpoch());

                        return new ProducerState(
                            activeProducer.producerId(),
                            activeProducer.producerEpoch(),
                            activeProducer.lastSequence(),
                            activeProducer.lastTimestamp(),
                            coordinatorEpoch,
                            currentTransactionFirstOffset
                        );
                    }).collect(Collectors.toList());

                completed.put(topicPartition, new PartitionProducerState(activeProducers));
            }
        }
        return new ApiResult<>(completed, failed, unmapped);
    }

}
