/*
 * 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.clients.consumer.AcknowledgeType;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.KafkaShareConsumer;
import org.apache.kafka.common.TopicIdPartition;
import org.apache.kafka.common.TopicPartition;

import java.time.Duration;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * {@link ShareFetch} represents the records fetched from the broker to be returned to the consumer
 * to satisfy a {@link KafkaShareConsumer#poll(Duration)} call. The records can come from multiple
 * topic-partitions.
 *
 * @param <K> The record key
 * @param <V> The record value
 */
public class ShareFetch<K, V> {
    private final Map<TopicIdPartition, ShareInFlightBatch<K, V>> batches;

    public static <K, V> ShareFetch<K, V> empty() {
        return new ShareFetch<>(new HashMap<>());
    }

    private ShareFetch(Map<TopicIdPartition, ShareInFlightBatch<K, V>> batches) {
        this.batches = batches;
    }

    /**
     * Add another {@link ShareInFlightBatch} to this one; all of its records will be added to this object's
     * {@link #records() records}.
     *
     * @param partition the topic-partition
     * @param batch the batch to add; may not be null
     */
    public void add(TopicIdPartition partition, ShareInFlightBatch<K, V> batch) {
        Objects.requireNonNull(batch);
        ShareInFlightBatch<K, V> currentBatch = this.batches.get(partition);
        if (currentBatch == null) {
            this.batches.put(partition, batch);
        } else {
            // This case shouldn't usually happen because we only send one fetch at a time per partition,
            // but it might conceivably happen in some rare cases (such as partition leader changes).
            currentBatch.merge(batch);
        }
    }

    /**
     * @return all the non-control messages for this fetch, grouped by partition
     */
    public Map<TopicPartition, List<ConsumerRecord<K, V>>> records() {
        final LinkedHashMap<TopicPartition, List<ConsumerRecord<K, V>>> result = new LinkedHashMap<>();
        batches.forEach((tip, batch) -> result.put(tip.topicPartition(), batch.getInFlightRecords()));
        return Collections.unmodifiableMap(result);
    }

    /**
     * @return the total number of non-control messages for this fetch, across all partitions
     */
    public int numRecords() {
        int numRecords = 0;
        if (!batches.isEmpty()) {
            Iterator<Map.Entry<TopicIdPartition, ShareInFlightBatch<K, V>>> iterator = batches.entrySet().iterator();
            while (iterator.hasNext()) {
                Map.Entry<TopicIdPartition, ShareInFlightBatch<K, V>> entry = iterator.next();
                ShareInFlightBatch<K, V> batch = entry.getValue();
                if (batch.isEmpty()) {
                    iterator.remove();
                } else {
                    numRecords += batch.numRecords();
                }
            }
        }

        return numRecords;
    }

    /**
     * @return {@code true} if and only if this fetch did not return any non-control records
     */
    public boolean isEmpty() {
        return numRecords() == 0;
    }

    /**
     * Acknowledge a single record in the current batch.
     *
     * @param record The record to acknowledge
     * @param type The acknowledge type which indicates whether it was processed successfully
     */
    public void acknowledge(final ConsumerRecord<K, V> record, AcknowledgeType type) {
        for (Map.Entry<TopicIdPartition, ShareInFlightBatch<K, V>> tipBatch : batches.entrySet()) {
            TopicIdPartition tip = tipBatch.getKey();
            if (tip.topic().equals(record.topic()) && (tip.partition() == record.partition())) {
                tipBatch.getValue().acknowledge(record, type);
                return;
            }
        }
        throw new IllegalStateException("The record cannot be acknowledged.");
    }

    /**
     * Acknowledge all records in the current batch. If any records in the batch already have
     * been acknowledged, those acknowledgements are not overwritten.
     *
     * @param type The acknowledge type which indicates whether it was processed successfully
     */
    public void acknowledgeAll(final AcknowledgeType type) {
        batches.forEach((tip, batch) -> batch.acknowledgeAll(type));
    }

    /**
     * Removes all acknowledged records from the in-flight records and returns the map of acknowledgements
     * to send. If some records were not acknowledged, the in-flight records will not be empty after this
     * method.
     *
     * @return The map of acknowledgements to send
     */
    public Map<TopicIdPartition, Acknowledgements> takeAcknowledgedRecords() {
        Map<TopicIdPartition, Acknowledgements> acknowledgementMap = new LinkedHashMap<>();
        batches.forEach((tip, batch) -> {
            Acknowledgements acknowledgements = batch.takeAcknowledgedRecords();
            if (!acknowledgements.isEmpty())
                acknowledgementMap.put(tip, acknowledgements);
        });
        return acknowledgementMap;
    }
}
