/*
 * 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.common.KafkaException;
import org.apache.kafka.common.protocol.Errors;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * This class maintains the acknowledgement and gap information for a set of records on a single
 * topic-partition being delivered to a consumer in a share group.
 */
public class Acknowledgements {
    public static final byte ACKNOWLEDGE_TYPE_GAP = (byte) 0;
    public static final int MAX_RECORDS_WITH_SAME_ACKNOWLEDGE_TYPE = 10;

    // The acknowledgements keyed by offset. If the record is a gap, the AcknowledgeType will be null.
    private final Map<Long, AcknowledgeType> acknowledgements;

    // When the broker responds to the acknowledgements, this is the exception thrown.
    private KafkaException acknowledgeException;

    // Set when the broker has responded to the acknowledgements.
    private boolean completed;

    public static Acknowledgements empty() {
        return new Acknowledgements(new TreeMap<>());
    }

    private Acknowledgements(Map<Long, AcknowledgeType> acknowledgements) {
        this.acknowledgements = acknowledgements;
        this.acknowledgeException = null;
        this.completed = false;
    }

    /**
     * Adds an acknowledgement for a specific offset. Will overwrite an existing
     * acknowledgement for the same offset.
     *
     * @param offset The record offset.
     * @param type   The AcknowledgeType.
     */
    public void add(long offset, AcknowledgeType type) {
        this.acknowledgements.put(offset, type);
    }

    /**
     * Adds an acknowledgement for a specific offset. Will <b>not</b> overwrite an existing
     * acknowledgement for the same offset.
     *
     * @param offset The record offset.
     * @param type   The AcknowledgeType.
     *
     * @return Whether the acknowledgement was added.
     */
    public boolean addIfAbsent(long offset, AcknowledgeType type) {
        return acknowledgements.putIfAbsent(offset, type) == null;
    }

    /**
     * Adds a gap for the specified offset. This means the broker expected a record at this offset
     * but when the record batch was parsed, the desired offset was missing. This occurs when the
     * record has been removed by the log compactor.
     *
     * @param offset The record offset.
     */
    public void addGap(long offset) {
        acknowledgements.put(offset, null);
    }

    /**
     * Gets the acknowledgement type for an offset.
     *
     * @param offset The record offset.
     * @return The AcknowledgeType, or null if no acknowledgement is present.
     */
    public AcknowledgeType get(long offset) {
        return acknowledgements.get(offset);
    }

    /**
     * Whether the set of acknowledgements is empty.
     *
     * @return Whether the set of acknowledgements is empty.
     */
    public boolean isEmpty() {
        return acknowledgements.isEmpty();
    }

    /**
     * Returns the size of the set of acknowledgements.
     *
     * @return The size of the set of acknowledgements.
     */
    public int size() {
        return acknowledgements.size();
    }

    /**
     * Whether the acknowledgements were sent to the broker and a response received.
     *
     * @return Whether the acknowledgements were sent to the broker and a response received
     */
    public boolean isCompleted() {
        return completed;
    }

    /**
     * Completes the acknowledgements when the response has been received from the broker.
     *
     * @param acknowledgeException the exception (will be null if successful)
     */
    public void complete(KafkaException acknowledgeException) {
        this.acknowledgeException = acknowledgeException;
        completed = true;
    }

    /**
     * Get the acknowledgement exception when the response has been received from the broker.
     *
     * @return the error code
     */
    public KafkaException getAcknowledgeException() {
        return acknowledgeException;
    }

    /**
     * Whether an acknowledgement error code was received in the response from the broker.
     *
     * @return Whether an acknowledgement error code was received in the response from the broker.
     */
    public boolean isCompletedExceptionally() {
        return acknowledgeException != null;
    }

    /**
     * Merges two sets of acknowledgements. If there are overlapping acknowledgements, the
     * merged set wins.
     *
     * @param other The set of acknowledgement to merge.
     */
    public Acknowledgements merge(Acknowledgements other) {
        acknowledgements.putAll(other.acknowledgements);
        return this;
    }

    /**
     * Returns the Map of Acknowledgements for the offsets.
     */
    public Map<Long, AcknowledgeType> getAcknowledgementsTypeMap() {
        return acknowledgements;
    }

    /**
     * Converts the acknowledgements into a list of {@link AcknowledgementBatch} which can easily
     * be converted into the form required for the RPC requests.
     */
    public List<AcknowledgementBatch> getAcknowledgementBatches() {
        List<AcknowledgementBatch> batches = new ArrayList<>();
        if (acknowledgements.isEmpty())
            return batches;

        AcknowledgementBatch currentBatch = null;
        for (Map.Entry<Long, AcknowledgeType> entry : acknowledgements.entrySet()) {
            if (currentBatch == null) {
                currentBatch = new AcknowledgementBatch();
                currentBatch.setFirstOffset(entry.getKey());
            } else {
                currentBatch = maybeCreateNewBatch(currentBatch, entry.getKey(), batches);
            }
            currentBatch.setLastOffset(entry.getKey());
            if (entry.getValue() != null) {
                currentBatch.acknowledgeTypes().add(entry.getValue().id);
            } else {
                currentBatch.acknowledgeTypes().add(ACKNOWLEDGE_TYPE_GAP);
            }
        }
        List<AcknowledgementBatch> optimalBatches = maybeOptimiseAcknowledgeTypes(currentBatch);

        optimalBatches.forEach(batch -> {
            if (canOptimiseForSingleAcknowledgeType(batch)) {
                // If the batch had a single acknowledgement type, we optimise the array independent
                // of the number of records.
                batch.acknowledgeTypes().subList(1, batch.acknowledgeTypes().size()).clear();
            }
            batches.add(batch);
        });
        return batches;
    }

    /**
     * Creates a new current batch if the next offset is not one higher than the current batch's
     * last offset.
     */
    private AcknowledgementBatch maybeCreateNewBatch(AcknowledgementBatch currentBatch, Long nextOffset, List<AcknowledgementBatch> batches) {
        if (nextOffset != currentBatch.lastOffset() + 1) {
            List<AcknowledgementBatch> optimalBatches = maybeOptimiseAcknowledgeTypes(currentBatch);

            optimalBatches.forEach(batch -> {
                if (canOptimiseForSingleAcknowledgeType(batch)) {
                    // If the batch had a single acknowledgement type, we optimise the array independent
                    // of the number of records.
                    batch.acknowledgeTypes().subList(1, batch.acknowledgeTypes().size()).clear();
                }
                batches.add(batch);
            });

            currentBatch = new AcknowledgementBatch();
            currentBatch.setFirstOffset(nextOffset);
        }

        return currentBatch;
    }

    /**
     * Traverses the acknowledgementBatch and splits it into optimal batches wherever possible.
     * Optimisation happens when a batch has continuous records of the same acknowledge type
     * whose count exceeds the default value. In this case, the batch is split into 2 such that the
     * batch with the continuous records has only 1 acknowledge type in its array.
     */
    private List<AcknowledgementBatch> maybeOptimiseAcknowledgeTypes(AcknowledgementBatch currentAcknowledgeBatch) {
        List<AcknowledgementBatch> batches = new ArrayList<>();
        if (currentAcknowledgeBatch == null) return batches;

        long currentOffset = currentAcknowledgeBatch.firstOffset();
        int currentStartIndex = 0;
        int recordsWithSameAcknowledgeType = 1;
        for (int i = 1; i < currentAcknowledgeBatch.acknowledgeTypes().size(); i++) {
            byte acknowledgeType = currentAcknowledgeBatch.acknowledgeTypes().get(i);
            // If we have a continuous set of records with the same acknowledgement type exceeding the default count,
            // then we optimise the batches to include only start and end offset and have only 1 acknowledge type in the array.
            byte prevAcknowledgeType = currentAcknowledgeBatch.acknowledgeTypes().get(i - 1);
            if (acknowledgeType == prevAcknowledgeType && recordsWithSameAcknowledgeType >= MAX_RECORDS_WITH_SAME_ACKNOWLEDGE_TYPE) {
                // We continue traversing until we have the same acknowledge type.
                while (i < currentAcknowledgeBatch.acknowledgeTypes().size()) {
                    byte acknowledgeType2 = currentAcknowledgeBatch.acknowledgeTypes().get(i);

                    if (acknowledgeType2 != currentAcknowledgeBatch.acknowledgeTypes().get(i - 1)) break;
                    i++;
                    recordsWithSameAcknowledgeType++;
                }

                // Now we prepare 2 batches, one starting just before the batch with single acknowledge type
                // and one with the single acknowledge type.
                AcknowledgementBatch batch1 = new AcknowledgementBatch();
                batch1.setFirstOffset(currentOffset);
                batch1.setLastOffset(currentOffset + i - recordsWithSameAcknowledgeType - currentStartIndex - 1);
                if (batch1.lastOffset() >= batch1.firstOffset()) {
                    batch1.setAcknowledgeTypes(new ArrayList<>(currentAcknowledgeBatch.acknowledgeTypes().subList(currentStartIndex,
                            i - recordsWithSameAcknowledgeType)));
                    batches.add(batch1);
                }

                AcknowledgementBatch batch2 = new AcknowledgementBatch();
                batch2.setFirstOffset(currentOffset + i - recordsWithSameAcknowledgeType - currentStartIndex);
                batch2.setLastOffset(currentOffset + i - currentStartIndex - 1);
                batch2.acknowledgeTypes().add(acknowledgeType);

                batches.add(batch2);
                recordsWithSameAcknowledgeType = 1;

                // Updating the offset and startIndex for further iterations.
                currentOffset = currentOffset + i - currentStartIndex;
                currentStartIndex = i;
            } else if (acknowledgeType == prevAcknowledgeType) {
                // The maximum limit has not yet been reached, we increment the count and move ahead.
                recordsWithSameAcknowledgeType++;
            } else {
                recordsWithSameAcknowledgeType = 1;
            }
        }
        if (currentStartIndex < currentAcknowledgeBatch.acknowledgeTypes().size()) {
            AcknowledgementBatch batch = new AcknowledgementBatch();
            batch.setFirstOffset(currentOffset);
            batch.setLastOffset(currentOffset + currentAcknowledgeBatch.acknowledgeTypes().size() - currentStartIndex - 1);
            batch.setAcknowledgeTypes(new ArrayList<>(currentAcknowledgeBatch.acknowledgeTypes().subList(currentStartIndex,
                    currentAcknowledgeBatch.acknowledgeTypes().size())));
            batches.add(batch);
        }
        return batches;
    }

    /**
     * @return Returns true if the array of acknowledge types in the share fetch batch contains a single acknowledge type
     * and the array size can be reduced to 1.
     * Returns false when the array has more than one acknowledge type or is already optimised.
     */
    private boolean canOptimiseForSingleAcknowledgeType(AcknowledgementBatch acknowledgementBatch) {
        if (acknowledgementBatch == null || acknowledgementBatch.acknowledgeTypes().size() == 1) return false;
        int firstAcknowledgeType = acknowledgementBatch.acknowledgeTypes().get(0);
        for (int i = 1; i < acknowledgementBatch.acknowledgeTypes().size(); i++) {
            if (acknowledgementBatch.acknowledgeTypes().get(i) != firstAcknowledgeType) return false;
        }
        return true;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder("Acknowledgements(");
        sb.append(acknowledgements);
        sb.append(", acknowledgeException=");
        sb.append(acknowledgeException != null ? Errors.forException(acknowledgeException) : "null");
        sb.append(", completed=");
        sb.append(completed);
        sb.append(")");
        return sb.toString();
    }
}
