package com.gfg;

import org.apache.kafka.clients.admin.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;

import java.util.*;
import java.util.concurrent.ExecutionException;

/**
 * KafkaTopicMessageCount class is used to calculate the number of messages in a Kafka topic.
 * It connects to a Kafka broker, retrieves offsets for each partition of the topic, and calculates the message count.
 */
public class KafkaTopicMessageCount {

    /**
     * Main method to execute the message count calculation.
     * 
     * @param args command line arguments
     * @throws ExecutionException if the computation threw an exception
     * @throws InterruptedException if the current thread was interrupted while waiting
     */
    public static void main(String[] args) throws ExecutionException, InterruptedException {
        // Define the topic to be analyzed
        String topic = "my-new-topic";
        
        // Set up properties for the AdminClient
        Properties properties = new Properties();
        properties.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");

        // Create an AdminClient with the specified properties
        try (AdminClient adminClient = AdminClient.create(properties)) {
            // Get the message counts for the specified topic
            Map<TopicPartition, Long> messageCounts = getMessageCountsForTopic(adminClient, topic);
            
            // Print the message counts for each partition
            messageCounts.forEach((tp, count) -> System.out.println("Partition: " + tp.partition() + ", Message Count: " + count));
        }
    }

    /**
     * Retrieves the message counts for each partition of the specified topic.
     * 
     * @param adminClient the AdminClient instance to interact with the Kafka broker
     * @param topic the name of the Kafka topic
     * @return a map of TopicPartition to message count
     * @throws ExecutionException if the computation threw an exception
     * @throws InterruptedException if the current thread was interrupted while waiting
     */
    public static Map<TopicPartition, Long> getMessageCountsForTopic(AdminClient adminClient, String topic) throws ExecutionException, InterruptedException {
        Map<TopicPartition, Long> messageCounts = new HashMap<>();

        // Fetch the beginning offsets for each partition
        List<TopicPartition> partitions = getTopicPartitions(adminClient, topic);
        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> beginningOffsets = adminClient.listOffsets(
                partitions.stream().collect(HashMap::new, (m, v) -> m.put(v, OffsetSpec.earliest()), HashMap::putAll)
        ).all().get();

        // Fetch the end offsets for each partition
        Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> endOffsets = adminClient.listOffsets(
                partitions.stream().collect(HashMap::new, (m, v) -> m.put(v, OffsetSpec.latest()), HashMap::putAll)
        ).all().get();

        // Calculate the number of messages in each partition
        for (TopicPartition partition : partitions) {
            long beginningOffset = beginningOffsets.get(partition).offset();
            long endOffset = endOffsets.get(partition).offset();
            long messageCount = endOffset - beginningOffset;
            messageCounts.put(partition, messageCount);
        }

        return messageCounts;
    }

    /**
     * Retrieves the list of TopicPartitions for the specified topic.
     * 
     * @param adminClient the AdminClient instance to interact with the Kafka broker
     * @param topic the name of the Kafka topic
     * @return a list of TopicPartitions
     * @throws ExecutionException if the computation threw an exception
     * @throws InterruptedException if the current thread was interrupted while waiting
     */
    private static List<TopicPartition> getTopicPartitions(AdminClient adminClient, String topic) throws ExecutionException, InterruptedException {
        // Describe the topic to get information about its partitions
        DescribeTopicsResult describeTopicsResult = adminClient.describeTopics(Collections.singletonList(topic));
        TopicDescription topicDescription = describeTopicsResult.all().get().get(topic);
        List<TopicPartition> partitions = new ArrayList<>();
        
        // Create a list of TopicPartitions from the topic description
        for (TopicPartitionInfo partitionInfo : topicDescription.partitions()) {
            partitions.add(new TopicPartition(topic, partitionInfo.partition()));
        }
        return partitions;
    }
}