package com.bcxin.message.configs;

import com.bcxin.message.common.CommonConstant;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.kafka.annotation.EnableKafka;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.listener.CommonErrorHandler;
import org.springframework.kafka.listener.DeadLetterPublishingRecoverer;
import org.springframework.kafka.listener.DefaultErrorHandler;

import java.util.function.BiFunction;

@EnableKafka
@Configuration
public class KafkaConfig {
    private static final Logger logger = LoggerFactory.getLogger(KafkaConfig.class);

    @Bean
    public CommonErrorHandler commonErrorHandler(KafkaTemplate kafkaTemplate) {
        final BiFunction<ConsumerRecord<?, ?>, Exception, TopicPartition>
                destinationResolver = (cr, e) -> new TopicPartition(CommonConstant.TOPIC_DEAD_LETTER_DISPATCH, cr.partition());

        BiFunction<ConsumerRecord<?, ?>, Exception, Headers> DEFAULT_HEADERS_FUNCTION =
                (rec, ex) -> {
                    Headers topicNameHeader = new RecordHeaders();
                    topicNameHeader.add(CommonConstant.HEADER_ORIGINAL_TOPIC_NAME, rec.topic().getBytes());

                    return topicNameHeader;
                };
        DeadLetterPublishingRecoverer deadLetterPublishingRecoverer = new DeadLetterPublishingRecoverer(kafkaTemplate, destinationResolver) {
            @Override
            public void accept(ConsumerRecord<?, ?> record, Consumer<?, ?> consumer, Exception exception) {
                boolean disallowedOrException = true;
                try {
                    if (!CommonConstant.TOPIC_DEAD_LETTER_DISPATCH.equalsIgnoreCase(record.topic())) {
                        super.accept(record, consumer, exception);
                        disallowedOrException = false;
                    }
                } catch (Exception ex) {
                    disallowedOrException = true;
                }

                if (disallowedOrException) {
                    logger.error(String.format("send data to dead letter queue while error:topic=%s;key=%s;", record.topic(), record.key()));
                }
            }
        };

        deadLetterPublishingRecoverer.setHeadersFunction(DEFAULT_HEADERS_FUNCTION);
        return new DefaultErrorHandler(deadLetterPublishingRecoverer);
    }
}
