/*
 * 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;

import org.apache.kafka.clients.KafkaClient;
import org.apache.kafka.clients.MockClient;
import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy;
import org.apache.kafka.clients.consumer.internals.ShareConsumerMetadata;
import org.apache.kafka.clients.consumer.internals.SubscriptionState;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicIdPartition;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.compress.Compression;
import org.apache.kafka.common.internals.ClusterResourceListeners;
import org.apache.kafka.common.message.ShareAcknowledgeResponseData;
import org.apache.kafka.common.message.ShareFetchResponseData;
import org.apache.kafka.common.message.ShareGroupHeartbeatResponseData;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.record.MemoryRecords;
import org.apache.kafka.common.record.MemoryRecordsBuilder;
import org.apache.kafka.common.record.TimestampType;
import org.apache.kafka.common.requests.FindCoordinatorResponse;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.requests.RequestTestUtils;
import org.apache.kafka.common.requests.ShareAcknowledgeRequest;
import org.apache.kafka.common.requests.ShareAcknowledgeResponse;
import org.apache.kafka.common.requests.ShareFetchRequest;
import org.apache.kafka.common.requests.ShareFetchResponse;
import org.apache.kafka.common.requests.ShareGroupHeartbeatRequest;
import org.apache.kafka.common.requests.ShareGroupHeartbeatResponse;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;

import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

// This test exercises the KafkaShareConsumer with the MockClient to validate the Kafka protocol RPCs
@Timeout(value = 120)
@SuppressWarnings({"ClassDataAbstractionCoupling"})
public class KafkaShareConsumerTest {

    private final String groupId = "test-group";
    private final String clientId1 = "client-id-1";

    private final String topic1 = "test1";
    private final Uuid topicId1 = Uuid.randomUuid();
    private final TopicPartition t1p0 = new TopicPartition(topic1, 0);
    private final TopicIdPartition ti1p0 = new TopicIdPartition(topicId1, t1p0);

    private final Map<String, Uuid> topicIds = Map.of(topic1, topicId1);

    private final int batchSize = 10;
    private final int heartbeatIntervalMs = 1000;

    private final Time time = new MockTime();
    private final SubscriptionState subscription = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.EARLIEST);

    @Test
    public void testVerifyHeartbeats() throws InterruptedException {
        ShareConsumerMetadata metadata = new ShareConsumerMetadata(0, 0, Long.MAX_VALUE, false,
            subscription, new LogContext(), new ClusterResourceListeners());
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic1, 1));
        Node node = metadata.fetch().nodes().get(0);

        Node coordinator = findCoordinator(client, node);

        // The member ID is created by the consumer and then sent in the initial request. The responses
        // need to contain the same member ID.
        final AtomicReference<Uuid> memberId = new AtomicReference<>();
        final AtomicInteger heartbeatsReceived = new AtomicInteger();
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareGroupHeartbeatRequest) {
                ShareGroupHeartbeatRequest request = (ShareGroupHeartbeatRequest) body;
                memberId.set(Uuid.fromString(request.data().memberId()));
                boolean matches = request.data().memberEpoch() == 0;
                heartbeatsReceived.addAndGet(1);

                client.prepareResponseFrom(body2 -> {
                    if (body2 instanceof ShareGroupHeartbeatRequest) {
                        ShareGroupHeartbeatRequest request2 = (ShareGroupHeartbeatRequest) body2;
                        boolean matches2 = request2.data().memberId().equals(memberId.get().toString()) && request2.data().memberEpoch() == 1;
                        heartbeatsReceived.addAndGet(1);
                        return matches2;
                    } else {
                        return false;
                    }
                }, shareGroupHeartbeatResponse(memberId.get(), 2, ti1p0), coordinator);

                return matches;
            } else {
                return false;
            }
        }, shareGroupHeartbeatResponse(memberId.get(), 1, ti1p0), coordinator);

        try (KafkaShareConsumer<String, String> consumer = newShareConsumer(clientId1, metadata, client)) {
            consumer.subscribe(Set.of(topic1));
            consumer.poll(Duration.ZERO);

            Thread.sleep(heartbeatIntervalMs);

            assertEquals(2, heartbeatsReceived.get());
            assertTrue(client.futureResponses().isEmpty());

            consumer.close(Duration.ZERO);
        }
    }

    @Test
    public void testVerifyFetchAndCommitSyncImplicit() {
        ShareConsumerMetadata metadata = new ShareConsumerMetadata(0, 0, Long.MAX_VALUE, false,
            subscription, new LogContext(), new ClusterResourceListeners());
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic1, 1));
        Node node = metadata.fetch().nodes().get(0);

        Node coordinator = findCoordinator(client, node);

        final AtomicReference<Uuid> memberId = new AtomicReference<>();
        final AtomicBoolean memberLeft = shareGroupHeartbeatGenerator(client, coordinator, memberId, ti1p0);

        // [A] A SHARE_FETCH in a new share session, fetching from topic topicId1, with no acknowledgements included.
        // The response includes 2 records which are acquired.
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareFetchRequest) {
                ShareFetchRequest request = (ShareFetchRequest) body;
                return request.data().groupId().equals(groupId) &&
                    request.data().shareSessionEpoch() == 0 &&
                    request.data().batchSize() == batchSize &&
                    request.data().topics().stream().findFirst().get().topicId().equals(topicId1) &&
                    request.data().topics().stream().findFirst().get().partitions().size() == 1 &&
                    request.data().topics().stream().findFirst().get().partitions().stream().findFirst().get().acknowledgementBatches().isEmpty();
            } else {
                return false;
            }
        }, shareFetchResponse(ti1p0, 2), node);

        // [B] A SHARE_ACKNOWLEDGE for the two records acquired in [A].
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareAcknowledgeRequest) {
                ShareAcknowledgeRequest request = (ShareAcknowledgeRequest) body;
                return request.data().groupId().equals(groupId) &&
                    request.data().shareSessionEpoch() == 1 &&
                    request.data().topics().stream().findFirst().get().partitions().stream().findFirst().get().acknowledgementBatches().get(0).firstOffset() == 0 &&
                    request.data().topics().stream().findFirst().get().partitions().stream().findFirst().get().acknowledgementBatches().get(0).lastOffset() == 1 &&
                    request.data().topics().stream().findFirst().get().partitions().stream().findFirst().get().acknowledgementBatches().get(0).acknowledgeTypes().size() == 1 &&
                    request.data().topics().stream().findFirst().get().partitions().stream().findFirst().get().acknowledgementBatches().get(0).acknowledgeTypes().get(0) == (byte) 1;
            } else {
                return false;
            }
        }, shareAcknowledgeResponse(ti1p0), node);

        // [C] A SHARE_ACKNOWLEDGE which closes the share session.
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareAcknowledgeRequest) {
                ShareAcknowledgeRequest request = (ShareAcknowledgeRequest) body;
                return request.data().groupId().equals(groupId) &&
                    request.data().shareSessionEpoch() == -1 &&
                    request.data().topics().isEmpty();
            } else {
                return false;
            }
        }, shareAcknowledgeResponse(), node);

        try (KafkaShareConsumer<String, String> consumer = newShareConsumer(clientId1, metadata, client)) {
            consumer.subscribe(Set.of(topic1));

            // This will be a SHARE_GROUP_HEARTBEAT to establish the membership and then a SHARE_FETCH [A]
            consumer.poll(Duration.ofMillis(5000));

            // This will be a SHARE_ACKNOWLEDGE [B]
            consumer.commitSync();

            // This will be a SHARE_ACKNOWLEDGE [C] and a final SHARE_GROUP_HEARTBEAT to leave the group
            consumer.close(Duration.ZERO);

            assertTrue(memberLeft.get());
            assertTrue(client.futureResponses().isEmpty());
        }
    }

    @Test
    public void testVerifyFetchAndCloseImplicit() {
        ShareConsumerMetadata metadata = new ShareConsumerMetadata(0, 0, Long.MAX_VALUE, false,
            subscription, new LogContext(), new ClusterResourceListeners());
        MockClient client = new MockClient(time, metadata);

        initMetadata(client, Map.of(topic1, 1));
        Node node = metadata.fetch().nodes().get(0);

        Node coordinator = findCoordinator(client, node);

        final AtomicReference<Uuid> memberId = new AtomicReference<>();
        final AtomicBoolean memberLeft = shareGroupHeartbeatGenerator(client, coordinator, memberId, ti1p0);

        // [A] A SHARE_FETCH in a new share session, fetching from topic topicId1, with no acknowledgements included.
        // The response includes 2 records which are acquired.
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareFetchRequest) {
                ShareFetchRequest request = (ShareFetchRequest) body;
                return request.data().groupId().equals(groupId) &&
                    request.data().shareSessionEpoch() == 0 &&
                    request.data().batchSize() == batchSize &&
                    request.data().topics().stream().findFirst().get().topicId().equals(topicId1) &&
                    request.data().topics().stream().findFirst().get().partitions().size() == 1 &&
                    request.data().topics().stream().findFirst().get().partitions().stream().findFirst().get().acknowledgementBatches().isEmpty();
            } else {
                return false;
            }
        }, shareFetchResponse(ti1p0, 2), node);

        // [B] A SHARE_ACKNOWLEDGE which closes the share session. Because this is implicit acknowledgement,
        // the acquired records are released by the broker when the share session is closed.
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareAcknowledgeRequest) {
                ShareAcknowledgeRequest request = (ShareAcknowledgeRequest) body;
                return request.data().groupId().equals(groupId) &&
                    request.data().shareSessionEpoch() == -1 &&
                    request.data().topics().isEmpty();
            } else {
                return false;
            }
        }, shareAcknowledgeResponse(), node);

        try (KafkaShareConsumer<String, String> consumer = newShareConsumer(clientId1, metadata, client)) {
            consumer.subscribe(Set.of(topic1));

            // This will be a SHARE_GROUP_HEARTBEAT to establish the membership and then a SHARE_FETCH [A]
            consumer.poll(Duration.ofMillis(5000));

            // This will be a SHARE_ACKNOWLEDGE [B] and a final SHARE_GROUP_HEARTBEAT to leave the group
            consumer.close(Duration.ZERO);

            assertTrue(memberLeft.get());
            assertTrue(client.futureResponses().isEmpty());
        }
    }

    private KafkaShareConsumer<String, String> newShareConsumer(String clientId,
                                                                ShareConsumerMetadata metadata,
                                                                KafkaClient client) {
        LogContext logContext = new LogContext();
        Deserializer<String> keyDeserializer = new StringDeserializer();
        Deserializer<String> valueDeserializer = new StringDeserializer();
        ShareConsumerConfig config = newConsumerConfig(clientId);

        return new KafkaShareConsumer<>(
            logContext,
            clientId,
            groupId,
            config,
            keyDeserializer,
            valueDeserializer,
            time,
            client,
            subscription,
            metadata
        );
    }

    private ShareConsumerConfig newConsumerConfig(String clientId) {
        Map<String, Object> configs = new HashMap<>();
        configs.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
        configs.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        configs.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, batchSize);
        configs.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
        return new ShareConsumerConfig(configs);
    }

    private void initMetadata(MockClient client, Map<String, Integer> partitions) {
        // It is important that this MetadataResponse contains the topic IDs for the topics which are being
        // subscribed to, otherwise the ShareMembershipManager will cause a second MetadataRequest to be sent out
        Map<String, Uuid> metadataTopicIds = new HashMap<>();
        for (String name : partitions.keySet()) {
            metadataTopicIds.put(name, topicIds.get(name));
        }
        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWithIds(1, partitions, metadataTopicIds);
        client.updateMetadata(initialMetadata);
    }

    private Node findCoordinator(MockClient client, Node node) {
        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
        return new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
    }

    // This method generates a sequence of prepared SHARE_GROUP_HEARTBEAT responses with increasing member epochs.
    // Each time that a SHARE_GROUP_HEARTBEAT response matches the prepared response matcher, the next prepared
    // response is added, until the matching requests member epoch is -1, indicating that the member is leaving
    // the group.
    private AtomicBoolean shareGroupHeartbeatGenerator(MockClient client, Node coordinator, AtomicReference<Uuid> memberId, TopicIdPartition tip) {
        AtomicBoolean memberLeft = new AtomicBoolean();
        AtomicInteger heartbeatsReceived = new AtomicInteger();
        shareGroupHeartbeat(client, coordinator, memberId, 0, tip, heartbeatsReceived, memberLeft);
        return memberLeft;
    }

    private void shareGroupHeartbeat(MockClient client, Node coordinator, AtomicReference<Uuid> memberId, int memberEpoch, TopicIdPartition tip, AtomicInteger heartbeatsReceived, AtomicBoolean memberLeft) {
        client.prepareResponseFrom(body -> {
            if (body instanceof ShareGroupHeartbeatRequest) {
                ShareGroupHeartbeatRequest request = (ShareGroupHeartbeatRequest) body;
                if (request.data().memberEpoch() == 0) {
                    memberId.set(Uuid.fromString(request.data().memberId()));
                }
                if (request.data().memberEpoch() == -1) {
                    memberLeft.set(true);
                } else {
                    shareGroupHeartbeat(client, coordinator, memberId, memberEpoch + 1, tip, heartbeatsReceived, memberLeft);
                }
                heartbeatsReceived.addAndGet(1);
                return true;
            } else {
                return false;
            }
        }, shareGroupHeartbeatResponse(memberId.get(), memberEpoch, tip), coordinator);
    }

    private ShareGroupHeartbeatResponse shareGroupHeartbeatResponse(Uuid memberId, int memberEpoch, TopicIdPartition tip) {
        if (memberEpoch != -1) {
            List<ShareGroupHeartbeatResponseData.TopicPartitions> assignedPartitions = new LinkedList<>();
            assignedPartitions.add(new ShareGroupHeartbeatResponseData.TopicPartitions().setTopicId(tip.topicId()).setPartitions(List.of(tip.partition())));

            return new ShareGroupHeartbeatResponse(
                new ShareGroupHeartbeatResponseData()
                    .setMemberId(memberId != null ? memberId.toString() : null)
                    .setMemberEpoch(memberEpoch)
                    .setHeartbeatIntervalMs(heartbeatIntervalMs)
                    .setAssignment(new ShareGroupHeartbeatResponseData.Assignment()
                        .setTopicPartitions(assignedPartitions))
            );
        } else {
            return new ShareGroupHeartbeatResponse(
                new ShareGroupHeartbeatResponseData()
                    .setMemberId(memberId != null ? memberId.toString() : null)
                    .setMemberEpoch(memberEpoch)
                    .setHeartbeatIntervalMs(heartbeatIntervalMs)
            );
        }
    }

    private ShareFetchResponse shareFetchResponse(TopicIdPartition tip, int count) {
        MemoryRecords records;
        try (MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), Compression.NONE,
            TimestampType.CREATE_TIME, 0)) {
            for (int i = 0; i < count; i++) {
                builder.append(0L, ("key-" + i).getBytes(), ("value-" + i).getBytes());
            }
            records = builder.build();
        }
        ShareFetchResponseData.PartitionData partData = new ShareFetchResponseData.PartitionData()
            .setPartitionIndex(tip.partition())
            .setRecords(records)
            .setAcquiredRecords(List.of(new ShareFetchResponseData.AcquiredRecords().setFirstOffset(0).setLastOffset(count - 1).setDeliveryCount((short) 1)));
        return ShareFetchResponse.of(Errors.NONE, 0, new LinkedHashMap<>(Map.of(tip, partData)), List.of(), 0);
    }

    private ShareAcknowledgeResponse shareAcknowledgeResponse() {
        return new ShareAcknowledgeResponse(
            new ShareAcknowledgeResponseData()
        );
    }

    private ShareAcknowledgeResponse shareAcknowledgeResponse(TopicIdPartition tip) {
        ShareAcknowledgeResponseData.PartitionData partData = new ShareAcknowledgeResponseData.PartitionData()
            .setPartitionIndex(tip.partition())
            .setErrorCode(Errors.NONE.code());
        ShareAcknowledgeResponseData.ShareAcknowledgeTopicResponse topicResponse = new ShareAcknowledgeResponseData.ShareAcknowledgeTopicResponse()
            .setTopicId(tip.topicId())
            .setPartitions(List.of(partData));
        return new ShareAcknowledgeResponse(
            new ShareAcknowledgeResponseData()
                .setResponses(new ShareAcknowledgeResponseData.ShareAcknowledgeTopicResponseCollection(List.of(topicResponse).iterator()))
        );
    }
}
