package co.dhan.tracing;

import brave.Tracing;
import brave.context.slf4j.MDCScopeDecorator;
import brave.propagation.CurrentTraceContext;
import brave.spring.rabbit.SpringRabbitTracing;
import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
import org.springframework.boot.autoconfigure.amqp.RabbitTemplateCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class RabbitTracingConfig {

    @Bean
    public CurrentTraceContext.ScopeDecorator logScopeDecorator() {
        return MDCScopeDecorator.newBuilder().build(); // Enables MDC injection of trace/span IDs
    }

    @Bean
    public SpringRabbitTracing springRabbitTracing(Tracing tracing) {
        return SpringRabbitTracing.newBuilder(tracing)
                .remoteServiceName("rabbitmq")
                .build();
    }

    @Bean
    public RabbitTemplateCustomizer tracingRabbitTemplateCustomizer(SpringRabbitTracing springRabbitTracing) {
        return rabbitTemplate -> springRabbitTracing.decorateRabbitTemplate(rabbitTemplate);
    }

    @Bean
    public SimpleRabbitListenerContainerFactory tracingListenerContainerFactory(
            SpringRabbitTracing springRabbitTracing,
            SimpleRabbitListenerContainerFactory factory) {

        // Decorate the factory AFTER Spring Boot has auto-configured it
        springRabbitTracing.decorateSimpleRabbitListenerContainerFactory(factory);
        return factory;
    }
}