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.MQDeferException;
import tech.easyflow.common.mq.core.MQMessage;
import tech.easyflow.common.mq.core.MQSubscription;

import java.util.List;

/**
 * 文档分块任务消费者。
 *
 * @author Codex
 * @since 2026-08-02
 */
@Component
public class DocumentImportSplitTaskConsumer implements MQConsumerHandler {

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

    private final KnowledgeDocumentImportTaskAppService appService;
    private final DocumentImportBatchCircuitBreaker circuitBreaker;
    private final MQProperties mqProperties;

    /**
     * 创建分块任务消费者。
     *
     * @param appService 文档导入应用服务
     * @param circuitBreaker 自动导入批次熔断器
     * @param mqProperties MQ 配置
     */
    public DocumentImportSplitTaskConsumer(
        KnowledgeDocumentImportTaskAppService appService,
        DocumentImportBatchCircuitBreaker circuitBreaker,
        MQProperties mqProperties) {
        this.appService = appService;
        this.circuitBreaker = circuitBreaker;
        this.mqProperties = mqProperties;
    }

    /**
     * 获取分块任务订阅。
     *
     * @return MQ 订阅信息
     */
    @Override
    public MQSubscription subscription() {
        MQSubscription subscription = new MQSubscription();
        subscription.setTopic(DocumentImportTaskMqConstants.SPLIT_TOPIC);
        subscription.setConsumerGroup(DocumentImportTaskMqConstants.SPLIT_GROUP);
        subscription.setShardCount(resolveShardCount());
        subscription.setBatchEnabled(false);
        return subscription;
    }

    /**
     * 处理一批分块任务消息。
     *
     * @param messages MQ 消息
     */
    @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;
            }
            try {
                appService.handleSplitTask(event.getTaskId());
            } catch (Exception error) {
                handleInfrastructureFailure(event, message, error);
            }
        }
    }

    /**
     * 记录完整异常并中断自动导入批次；中断状态暂时无法持久化时保留消息 pending。
     *
     * @param event 任务事件
     * @param message MQ 消息
     * @param error 原始异常
     */
    private void handleInfrastructureFailure(DocumentImportTaskMessage event,
                                             MQMessage message,
                                             Exception error) {
        try {
            if (circuitBreaker.interruptTaskBatch(event.getTaskId(), error)) {
                return;
            }
        } catch (Exception interruptError) {
            interruptError.addSuppressed(error);
            LOG.error(
                "文档分块异常且批次中断状态写入失败，消息将暂缓确认: "
                    + "taskId={}, messageId={}, streamMessageId={}",
                event.getTaskId(), message.getMessageId(),
                message.getStreamMessageId(), interruptError);
            throw new MQDeferException(
                "文档分块批次中断状态暂时无法持久化", interruptError);
        }
        LOG.error(
            "文档分块任务发生基础设施异常，消息将暂缓确认: "
                + "taskId={}, messageId={}, streamMessageId={}",
            event.getTaskId(), message.getMessageId(),
            message.getStreamMessageId(), error);
        throw new MQDeferException("文档分块任务暂时无法处理", error);
    }

    /**
     * 获取当前 Redis Stream 分片数。
     *
     * @return 分片数
     */
    private int resolveShardCount() {
        return Math.max(
            mqProperties.getRedis().getChatPersistShardCount(), 1);
    }
}
