/*
 * 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.pulsar.broker.service;

import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
import com.google.common.collect.Sets;
import io.netty.buffer.ByteBuf;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import lombok.RequiredArgsConstructor;
import org.apache.bookkeeper.mledger.impl.EntryImpl;
import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
import org.apache.pulsar.common.api.proto.MessageMetadata;
import org.apache.pulsar.common.protocol.Commands;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

@Test(groups = "broker")
public class SharedConsumerAssignorTest {

    private final ConsumerSelector roundRobinConsumerSelector = new ConsumerSelector();
    private final List<EntryAndMetadata> entryAndMetadataList = new ArrayList<>();
    private final List<EntryAndMetadata> replayQueue = new ArrayList<>();
    private final Queue<EntryAndMetadata> cleanupQueue = new LinkedBlockingDeque<>();

    private SharedConsumerAssignor assignor;

    @BeforeMethod
    public void prepareData() {
        roundRobinConsumerSelector.clear();
        entryAndMetadataList.clear();
        replayQueue.clear();
        assignor = new SharedConsumerAssignor(roundRobinConsumerSelector, replayQueue::add, null);
        final AtomicLong entryId = new AtomicLong(0L);
        final MockProducer producerA = new MockProducer("A", entryId, entryAndMetadataList);
        final MockProducer producerB = new MockProducer("B", entryId, entryAndMetadataList);
        producerA.sendMessage();
        producerA.sendChunk(0, 3);
        producerA.sendChunk(1, 3);
        producerB.sendMessage();
        producerB.sendChunk(0, 2);
        producerA.sendChunk(2, 3);
        producerB.sendChunk(1, 2);
        // Use the following data for all tests
        // | entry id | uuid | chunk id | number of chunks |
        // | -------- | ---- | -------- | ---------------- |
        // | 0        | A-0  |          |                  |
        // | 1        | A-1  | 0        | 3                |
        // | 2        | A-1  | 1        | 3                |
        // | 3        | B-0  |          |                  |
        // | 4        | B-1  | 0        | 2                |
        // | 5        | A-1  | 2        | 3                |
        // | 6        | B-1  | 1        | 2               |
        // P.S. In the table above, The uuid represents the "<producer-name>-<sequence-id>" for non-chunks
        assertEquals(toString(entryAndMetadataList), Arrays.asList(
                "0:0@A-0", "0:1@A-1-0-3", "0:2@A-1-1-3", "0:3@B-0", "0:4@B-1-0-2", "0:5@A-1-2-3", "0:6@B-1-1-2"));
        entryAndMetadataList.forEach(entry -> assertEquals(entry.getDataBuffer().refCnt(), 1));
        cleanupQueue.addAll(entryAndMetadataList);
    }

    @AfterMethod
    public void releaseEntries() {
        EntryAndMetadata entry;
        while ((entry = cleanupQueue.poll()) != null) {
            entry.release();
        }
    }

    @Test
    public void testSingleConsumerMultiAssign() {
        // Only first 5 entries can be received because the number of permits is 5
        final Consumer consumer = new Consumer("A", 5);
        roundRobinConsumerSelector.addConsumers(consumer);

        Map<Consumer, List<EntryAndMetadata>> result = assignor.assign(entryAndMetadataList, 1);
        assertEquals(result.getOrDefault(consumer, Collections.emptyList()), entryAndMetadataList.subList(0, 5));
        // Since two chunked messages (A-1 and B-1) are both not received, these uuids have been cached
        assertEquals(assignor.getUuidToConsumer().keySet(), Sets.newHashSet("A-1", "B-1"));
        assertEquals(toString(replayQueue), Arrays.asList("0:5@A-1-2-3", "0:6@B-1-1-2"));

        result = assignor.assign(entryAndMetadataList.subList(5, 6), 1);
        assertEquals(result.getOrDefault(consumer, Collections.emptyList()), entryAndMetadataList.subList(5, 6));
        // A-1 is received so that uuid "A-1" has been removed from the cache
        assertEquals(assignor.getUuidToConsumer().keySet(), Sets.newHashSet("B-1"));

        result = assignor.assign(entryAndMetadataList.subList(6, 7), 1);
        assertEquals(result.getOrDefault(consumer, Collections.emptyList()), entryAndMetadataList.subList(6, 7));
        assertTrue(assignor.getUuidToConsumer().isEmpty());
    }

    @Test
    public void testMultiConsumerWithSmallPermits() {
        final Consumer consumerA = new Consumer("A", 3);
        final Consumer consumerB = new Consumer("B", 4);
        roundRobinConsumerSelector.addConsumers(consumerA, consumerB);

        // The original order:
        //   A-0, A-1-0-3, A-1-1-3, B-0, B-1-0-2, A-1-2-3, B-1-1-2
        // First consumerA received 3 entries and the available permits became 0. Then the consumer is switched to
        // consumerB and then B-0 and B-1-0-2 were received by it.
        // When consumerB tried to receive A-1-2-3, since the uuid "A-1" was already assigned to consumerA, consumerB
        // is not able to receive A-1-2-3.
        // However, since consumerA has no more permits, A-1-2-3 cannot be assigned to consumerA as well.
        // Therefore, A-1-2-3 was skipped and added to the replay queue.
        Map<Consumer, List<EntryAndMetadata>> result = assignor.assign(entryAndMetadataList, 2);
        assertEquals(toString(result.getOrDefault(consumerA, Collections.emptyList())),
                Arrays.asList("0:0@A-0", "0:1@A-1-0-3", "0:2@A-1-1-3"));
        assertEquals(toString(result.getOrDefault(consumerB, Collections.emptyList())),
                Arrays.asList("0:3@B-0", "0:4@B-1-0-2", "0:6@B-1-1-2"));
        assertEquals(toString(replayQueue), Collections.singletonList("0:5@A-1-2-3"));
        assertEquals(assignor.getUuidToConsumer().keySet(), Sets.newHashSet("A-1"));

        roundRobinConsumerSelector.clear();
        roundRobinConsumerSelector.addConsumers(consumerB, consumerA);
        assertSame(roundRobinConsumerSelector.peek(), consumerB);
        final List<EntryAndMetadata> leftEntries = new ArrayList<>(replayQueue);
        replayQueue.clear();
        result = assignor.assign(leftEntries, 2);
        // Since uuid "A-1" is still cached, A-1-2-3 will be dispatched to consumerA even if consumerB is the first
        // consumer returned by roundRobinConsumerSelector.get()
        assertEquals(toString(result.getOrDefault(consumerA, Collections.emptyList())),
                Collections.singletonList("0:5@A-1-2-3"));
        assertNull(result.get(consumerB));
        assertTrue(replayQueue.isEmpty());
        assertTrue(assignor.getUuidToConsumer().isEmpty());
    }

    @RequiredArgsConstructor
    static class ConsumerSelector implements Supplier<Consumer> {

        private final List<Consumer> consumers = new ArrayList<>();
        private int index = 0;

        public void addConsumers(Consumer... consumers) {
            this.consumers.addAll(Arrays.asList(consumers));
        }

        public void clear() {
            consumers.clear();
            index = 0;
        }

        @Override
        public Consumer get() {
            // a simple round-robin dispatcher
            final Consumer consumer = peek();
            if (consumer == null) {
                return null;
            }
            index++;
            return consumer;
        }

        public Consumer peek() {
            if (consumers.isEmpty()) {
                return null;
            }
            return consumers.get(index % consumers.size());
        }
    }

    private static List<String> toString(final List<EntryAndMetadata> entryAndMetadataList) {
        return entryAndMetadataList.stream().map(EntryAndMetadata::toString).collect(Collectors.toList());
    }

    @RequiredArgsConstructor
    static class MockProducer {
        final String name;
        final AtomicLong entryId;
        final List<EntryAndMetadata> entryAndMetadataList;
        long sequenceId = 0L;

        void sendMessage() {
            entryAndMetadataList.add(createEntryAndMetadata(entryId.getAndIncrement(),
                    createMetadata(name, sequenceId++, null, null)));
        }

        void sendChunk(int chunkId, int numChunks) {
            entryAndMetadataList.add(createEntryAndMetadata(entryId.getAndIncrement(),
                    createMetadata(name, sequenceId, chunkId, numChunks)));
            if (chunkId == numChunks - 1) {
                sequenceId++;
            }
        }
    }

    private static EntryAndMetadata createEntryAndMetadata(final long entryId,
                                                           final MessageMetadata metadata) {
        ByteBuf payload = PulsarByteBufAllocator.DEFAULT.buffer();
        final ByteBuf data = Commands.serializeMetadataAndPayload(
                Commands.ChecksumType.Crc32c, metadata, payload);
        payload.release();
        EntryAndMetadata entryAndMetadata = EntryAndMetadata.create(EntryImpl.create(0L, entryId, data));
        data.release();
        return entryAndMetadata;
    }

    private static MessageMetadata createMetadata(final String producerName,
                                                  final long sequenceId,
                                                  final Integer chunkId,
                                                  final Integer numChunks) {
        final MessageMetadata metadata = new MessageMetadata();
        metadata.setProducerName(producerName);
        metadata.setSequenceId(sequenceId);
        metadata.setPublishTime(0L);
        if (chunkId != null && numChunks != null) {
            metadata.setUuid(producerName + "-" + sequenceId);
            metadata.setChunkId(chunkId);
            metadata.setNumChunksFromMsg(numChunks);
        }
        return metadata;
    }

    /**
     * When there are no consumers online, chunk messages will not be directly lost.
     */
    @Test
    public void testChunkMessagesNotBeLostNoConsumer() {
        // 1. No consumer initially
        Map<Consumer, List<EntryAndMetadata>> result = assignor.assign(entryAndMetadataList, 1);
        assertTrue(result.isEmpty());
        assertEquals(replayQueue.size(), entryAndMetadataList.size());
        assertEquals(toString(replayQueue), toString(entryAndMetadataList));

        // 2. Two Consumers come online
        final Consumer consumerA = new Consumer("A", 100);
        final Consumer consumerB = new Consumer("B", 100);
        roundRobinConsumerSelector.addConsumers(consumerA, consumerB);

        // 3. Retry messages from replay queue
        List<EntryAndMetadata> retryList = new ArrayList<>(replayQueue);
        replayQueue.clear();

        // Use a larger batch size to ensure we can process enough messages
        result = assignor.assign(retryList, 10);

        // 4. Verify consumer receives all messages
        int totalReceived = result.values().stream().mapToInt(List::size).sum();
        assertEquals(totalReceived, retryList.size());

        // Verify that chunks are assigned to the same consumer
        List<String> entriesA = toString(result.getOrDefault(consumerA, Collections.emptyList()));
        List<String> entriesB = toString(result.getOrDefault(consumerB, Collections.emptyList()));

        // Check A-1 chunks (0:1, 0:2, 0:5)
        boolean a1InA = entriesA.stream().anyMatch(s -> s.contains("A-1"));
        if (a1InA) {
            assertTrue(entriesA.containsAll(Arrays.asList("0:1@A-1-0-3", "0:2@A-1-1-3", "0:5@A-1-2-3")));
            assertTrue(entriesB.stream().noneMatch(s -> s.contains("A-1")));
        } else {
            assertTrue(entriesB.containsAll(Arrays.asList("0:1@A-1-0-3", "0:2@A-1-1-3", "0:5@A-1-2-3")));
            assertTrue(entriesA.stream().noneMatch(s -> s.contains("A-1")));
        }

        // Check B-1 chunks (0:4, 0:6)
        boolean b1InA = entriesA.stream().anyMatch(s -> s.contains("B-1"));
        if (b1InA) {
            assertTrue(entriesA.containsAll(Arrays.asList("0:4@B-1-0-2", "0:6@B-1-1-2")));
            assertTrue(entriesB.stream().noneMatch(s -> s.contains("B-1")));
        } else {
            assertTrue(entriesB.containsAll(Arrays.asList("0:4@B-1-0-2", "0:6@B-1-1-2")));
            assertTrue(entriesA.stream().noneMatch(s -> s.contains("B-1")));
        }

        // 5. Verify internal state is clean (since all chunks are completed)
        assertTrue(assignor.getUuidToConsumer().isEmpty());
    }

}
