package tech.easyflow.ai.documentimport.task;

import com.alibaba.fastjson2.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import tech.easyflow.common.mq.config.MQProperties;
import tech.easyflow.common.mq.core.MQConsumerHandler;
import tech.easyflow.common.mq.core.MQMessage;
import tech.easyflow.common.mq.core.MQSubscription;

import java.util.List;

/**
 * 文档向量化任务消费者。
 *
 * @author Codex
 * @since 2026-04-14
 */
@Component
public class DocumentImportIndexTaskConsumer implements MQConsumerHandler {

    private static final Logger LOG = LoggerFactory.getLogger(DocumentImportIndexTaskConsumer.class);

    private final KnowledgeDocumentImportTaskAppService appService;
    private final MQProperties mqProperties;

    public DocumentImportIndexTaskConsumer(KnowledgeDocumentImportTaskAppService appService,
                                          MQProperties mqProperties) {
        this.appService = appService;
        this.mqProperties = mqProperties;
    }

    @Override
    public MQSubscription subscription() {
        MQSubscription subscription = new MQSubscription();
        subscription.setTopic(DocumentImportTaskMqConstants.INDEX_TOPIC);
        subscription.setConsumerGroup(DocumentImportTaskMqConstants.INDEX_GROUP);
        subscription.setShardCount(resolveShardCount());
        return subscription;
    }

    @Override
    public void handle(List<MQMessage> messages) {
        LOG.info("文档向量化消费者收到消息批次: count={}", messages == null ? 0 : messages.size());
        for (MQMessage message : messages) {
            DocumentImportTaskMessage event = JSON.parseObject(message.getBody(), DocumentImportTaskMessage.class);
            if (event == null || event.getTaskId() == null) {
                LOG.warn("文档向量化消费者跳过非法消息: streamMessageId={}, messageId={}",
                    message == null ? null : message.getStreamMessageId(),
                    message == null ? null : message.getMessageId());
                continue;
            }
            LOG.info("文档向量化消费者开始处理消息: taskId={}, messageId={}, streamMessageId={}",
                event.getTaskId(), message.getMessageId(), message.getStreamMessageId());
            try {
                appService.handleIndexTask(event.getTaskId());
                LOG.info("文档向量化消费者处理完成: taskId={}, messageId={}, streamMessageId={}",
                    event.getTaskId(), message.getMessageId(), message.getStreamMessageId());
            } catch (Exception exception) {
                LOG.error("文档向量化消费者处理失败: taskId={}, messageId={}, streamMessageId={}",
                    event.getTaskId(), message.getMessageId(), message.getStreamMessageId(), exception);
                throw exception;
            }
        }
    }

    /**
     * 向量化消费者需覆盖生产端的所有分片，避免消息落入未订阅分片。
     *
     * @return 当前 Redis Stream 分片数
     */
    private int resolveShardCount() {
        return Math.max(mqProperties.getRedis().getChatPersistShardCount(), 1);
    }
}
