package com.ktg.common.utils.rocketmq;

import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.message.MessageExt;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * RocketMQ消费者管理器
 * 负责管理多个RocketMQ消费者实例
 * 实现InitializingBean在Spring初始化完成后自动初始化消费者
 * 实现DisposableBean在Spring销毁时自动关闭消费者
 */
public class RocketMQConsumerManager implements InitializingBean, DisposableBean {
    // RocketMQ配置属性
    private final RocketMQProperties properties;

    // 消息处理器Map，key为消费者名称，value为消息处理函数
    private final Map<String, Consumer<MessageExt>> messageHandlers;

    // 消费者Map，key为消费者名称，value为消费者实例
    private final Map<String, DefaultMQPushConsumer> consumerMap = new ConcurrentHashMap<>();

    /**
     * 构造函数
     * @param properties RocketMQ配置属性
     * @param messageHandlers 消息处理器Map
     */
    public RocketMQConsumerManager(RocketMQProperties properties,
                                   Map<String, Consumer<MessageExt>> messageHandlers) {
        this.properties = properties;
        this.messageHandlers = messageHandlers;
    }

    /**
     * Spring Bean初始化完成后调用的方法
     * 初始化所有配置的消费者实例
     */
    @Override
    public void afterPropertiesSet() {
        // 遍历所有消费者配置
        properties.getConsumers().forEach((name, config) -> {
            // 检查消费者是否启用
            if (!config.isEnabled()) {
                return;
            }

            // 获取对应的消息处理器
            Consumer<MessageExt> messageHandler = messageHandlers.get(name);
            if (messageHandler == null) {
                throw new IllegalStateException("No message handler found for consumer: " + name);
            }

            // 创建消费者实例
            DefaultMQPushConsumer consumer = new DefaultMQPushConsumer(config.getGroup());

            // 设置NameServer地址
            consumer.setNamesrvAddr(properties.getNamesrvAddr());

            // 设置消费者参数
            consumer.setConsumeThreadMin(config.getConsumeThreadMin());
            consumer.setConsumeThreadMax(config.getConsumeThreadMax());
            consumer.setConsumeMessageBatchMaxSize(config.getConsumeMessageBatchMaxSize());

            // 设置实例名称（如果配置了）
            if (config.getInstanceName() != null) {
                consumer.setInstanceName(config.getInstanceName());
            }

            try {
                // 订阅Topic和Tags
                consumer.subscribe(config.getTopic(), config.getTags());

                // 注册消息监听器
                consumer.registerMessageListener(new MessageListenerConcurrently() {
                    @Override
                    public ConsumeConcurrentlyStatus consumeMessage(List<MessageExt> msgs,
                                                                    ConsumeConcurrentlyContext context) {
                        // 处理每条消息
                        for (MessageExt msg : msgs) {
                            try {
                                // 调用消息处理器处理消息
                                messageHandler.accept(msg);
                            } catch (Exception e) {
                                // 处理失败，稍后重试
                                return ConsumeConcurrentlyStatus.RECONSUME_LATER;
                            }
                        }
                        // 处理成功
                        return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
                    }
                });
                try {
                    // 启动消费者
                    consumer.start();
                    System.out.println("=============初始化消费者成功===============");
                } catch (Exception e) {
                    System.out.println("初始化消费者失败！");
                }

                // 将消费者加入Map
                consumerMap.put(name, consumer);
            } catch (MQClientException e) {
                throw new RuntimeException("Failed to start RocketMQ consumer: " + name, e);
            }
        });
    }

    /**
     * Spring Bean销毁时调用的方法
     * 关闭所有消费者实例
     */
    @Override
    public void destroy() {
        System.out.println("=================销毁消费者====================");
        consumerMap.values().forEach(DefaultMQPushConsumer::shutdown);
    }

    /**
     * 获取指定名称的消费者实例
     * @param name 消费者名称
     * @return 消费者实例
     */
    public DefaultMQPushConsumer getConsumer(String name) {
        DefaultMQPushConsumer consumer = consumerMap.get(name);
        if (consumer == null) {
            throw new IllegalArgumentException("RocketMQ consumer not found: " + name);
        }
        return consumer;
    }
}
