package com.tencent.tcb.operation.ckafka.plain;

import com.google.common.collect.Lists;
import com.tencent.tcb.operation.ckafka.JavaKafkaConfigurer;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.config.SaslConfigs;

public class KafkaPlainConsumerDemo {

    public static void main(String args[]) {
        // Set the path of the JAAS configuration file.
        JavaKafkaConfigurer.configureSaslPlain();

        //Load kafka.properties.
        Properties kafkaProperties = JavaKafkaConfigurer.getKafkaProperties();

        Properties props = new Properties();
        //Set the access point. Obtain the corresponding topic access point from the console.
        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaProperties.getProperty("bootstrap.servers"));

        // Access protocol.
        props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_PLAINTEXT");
        // Plain method.
        props.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
        //The maximum allowed interval between two polls.
        //If the consumer does not return a heartbeat message within the interval, the broker determines that the consumer is not alive. the broker removes the consumer from the consumer group and triggers rebalancing. The default value is 30s.
        props.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 30000);
        //Set the maximum number of messages that can be polled at a time.
        //Be cautious not to set this value too high. If too much data is polled and not consumed before the next poll, a load balancing will be triggered, causing a delay.
        props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 30);
        //Set the method for deserializing messages.
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
                "org.apache.kafka.common.serialization.StringDeserializer");
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
                "org.apache.kafka.common.serialization.StringDeserializer");
        //Specify the consumer group to which the current consumer instance belongs. Enter the group you created in the console.
        //Instances within the same consumer group consume messages in a load-balanced manner.
        props.put(ConsumerConfig.GROUP_ID_CONFIG, kafkaProperties.getProperty("group.id"));

        // Position of consumption offset. Note! If auto.offset.reset is set to "none", the consumer group will report an error of not finding the offset during the first consumption. In this case, the offset needs to be manually set in the catch block during the first consumption.
        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none");
        //Construct a consumer object, which essentially generates a consumer instance.
        KafkaConsumer<String, String> consumer = new KafkaConsumer<String, String>(props);
        //Set the topics for the consumer group to subscribe to, multiple topics can be subscribed.
        //If the GROUP_ID_CONFIG is the same, it is recommended to set the subscribed topics to be the same as well.
        List<String> subscribedTopics = new ArrayList<String>();
        //To subscribe to multiple topics, simply add them here.
        //You must create the topics in the console in advance.
        String topicStr = kafkaProperties.getProperty("topic");
        String[] topics = topicStr.split(",");
        for (String topic : topics) {
            subscribedTopics.add(topic.trim());
        }
        consumer.subscribe(subscribedTopics);
        //Consume messages in a loop.
        while (true) {
            try {
                ConsumerRecords<String, String> records = consumer.poll(1000);
                // The data must be consumed before the next poll, and the total time taken should not exceed SESSION_TIMEOUT_MS_CONFIG. It is recommended to use a separate thread pool to consume messages and return results asynchronously.
                for (ConsumerRecord<String, String> record : records) {
                    System.out.println(
                            String.format("Consume partition:%d offset:%d", record.partition(), record.offset()));
                }
            } catch (NoOffsetForPartitionException e) {
                System.out.println(e.getMessage());

                None: When auto.offset.reset is set to 'none', you need to catch exceptions and set the offset yourself. You can choose one of the following methods based on your business situation.
                // e.g. 1: Specify the offset. Here, you need to maintain the offset yourself for easy retries.
                Map<Integer, Long> partitionBeginOffsetMap = getPartitionOffset(consumer, topicStr, true);
                Map<Integer, Long> partitionEndOffsetMap = getPartitionOffset(consumer, topicStr, false);
                consumer.seek(new TopicPartition(topicStr, 0), 0);

                // e.g. 2: Start consuming from the beginning
                consumer.seekToBeginning(Lists.newArrayList(new TopicPartition(topicStr, 0)));

                //e.g 3: Designate the offset as the most recently available offset.
                consumer.seekToEnd(Lists.newArrayList(new TopicPartition(topicStr, 0)));

                // e.g. 4: Obtain the offset based on the timestamp, which means setting the offset according to the timestamp. For example, reset to the offset 10 minutes ago.
                Map<TopicPartition, Long> timestampsToSearch = new HashMap<>();
                Long value = Instant.now().minus(300, ChronoUnit.SECONDS).toEpochMilli();
                timestampsToSearch.put(new TopicPartition(topicStr, 0), value);
                Map<TopicPartition, OffsetAndTimestamp> topicPartitionOffsetAndTimestampMap = consumer
                        .offsetsForTimes(timestampsToSearch);
                for (Entry<TopicPartition, OffsetAndTimestamp> entry : topicPartitionOffsetAndTimestampMap
                        .entrySet()) {
                    TopicPartition topicPartition = entry.getKey();
                    OffsetAndTimestamp entryValue = entry.getValue();
                    consumer.seek(topicPartition, entryValue.offset()); // Specify the offset. You need to maintain the offset yourself for easy retries.
                }


            }
        }
    }

    /**
     Retrieve the earliest and most recent offset of the topic.
     * @param consumer
     * @param topicStr
     * @param beginOrEnd true begin; false end
     * @return
     */
    private static Map<Integer, Long> getPartitionOffset(KafkaConsumer<String, String> consumer, String topicStr,
            boolean beginOrEnd) {
        Collection<PartitionInfo> partitionInfos = consumer.partitionsFor(topicStr);
        List<TopicPartition> tp = new ArrayList<>();
        Map<Integer, Long> map = new HashMap<>();
        partitionInfos.forEach(str -> tp.add(new TopicPartition(topicStr, str.partition())));
        Map<TopicPartition, Long> topicPartitionLongMap;
        if (beginOrEnd) {
            topicPartitionLongMap = consumer.beginningOffsets(tp);
        } else {
            topicPartitionLongMap = consumer.endOffsets(tp);
        }
        topicPartitionLongMap.forEach((key, beginOffset) -> {
            int partition = key.partition();
            map.put(partition, beginOffset);
        });
        return map;
    }

}