package tech.easyflow.agent.runtime.document;

import com.alibaba.fastjson2.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;
import tech.easyflow.agent.config.AgentDocumentProperties;
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;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * Agent 文档轻量读取任务消费者。
 */
@Component
public class AgentDocumentReadTaskConsumer implements MQConsumerHandler {

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

    private final AgentDocumentReaderService readerService;
    private final AgentDocumentProperties properties;
    private final MQProperties mqProperties;
    private final ExecutorService executor;

    /**
     * 创建读取任务消费者。
     *
     * @param readerService 文档读取服务
     * @param properties 文档配置
     * @param mqProperties MQ 配置
     * @param executor 独立有界执行器
     */
    public AgentDocumentReadTaskConsumer(
            AgentDocumentReaderService readerService,
            AgentDocumentProperties properties,
            MQProperties mqProperties,
            @Qualifier("agentDocumentReaderExecutor") ExecutorService executor) {
        this.readerService = readerService;
        this.properties = properties;
        this.mqProperties = mqProperties;
        this.executor = executor;
    }

    /**
     * 声明读取主题订阅。
     *
     * @return MQ 订阅
     */
    @Override
    public MQSubscription subscription() {
        MQSubscription subscription = new MQSubscription();
        subscription.setTopic(AgentDocumentTaskMqConstants.READ_TOPIC);
        subscription.setConsumerGroup(AgentDocumentTaskMqConstants.READ_GROUP);
        subscription.setShardCount(Math.max(mqProperties.getRedis().getChatPersistShardCount(), 1));
        subscription.setBatchEnabled(false);
        return subscription;
    }

    /**
     * 将 MQ 事件提交到文档独立线程池并执行超时控制。
     *
     * @param messages MQ 消息
     * @throws Exception 读取失败或超时
     */
    @Override
    public void handle(List<MQMessage> messages) throws Exception {
        for (MQMessage message : messages == null ? List.<MQMessage>of() : messages) {
            AgentDocumentTaskMessage event = message == null ? null
                    : JSON.parseObject(message.getBody(), AgentDocumentTaskMessage.class);
            if (event == null || event.getAttachmentId() == null || event.getAttachmentId().isBlank()) {
                LOG.warn("跳过非法 Agent 文档读取消息: messageId={}",
                        message == null ? null : message.getMessageId());
                continue;
            }
            process(event.getAttachmentId());
        }
    }

    private void process(String attachmentId) throws Exception {
        Future<?> future;
        try {
            future = executor.submit(() -> readerService.process(attachmentId));
        } catch (RejectedExecutionException error) {
            readerService.markBusy(attachmentId);
            throw new IllegalStateException("Agent document reader queue is full", error);
        }
        try {
            future.get(properties.getReadTimeout().toMillis(), TimeUnit.MILLISECONDS);
        } catch (TimeoutException error) {
            // 先提交超时状态，避免工作线程收到中断后抢先写成普通“已取消”。
            readerService.markTimeout(attachmentId);
            future.cancel(true);
            throw new IllegalStateException("Agent document read timeout", error);
        } catch (InterruptedException error) {
            future.cancel(true);
            Thread.currentThread().interrupt();
            throw error;
        } catch (ExecutionException error) {
            Throwable cause = error.getCause();
            if (cause instanceof Exception exception) {
                throw exception;
            }
            throw new IllegalStateException("Agent document read failed", cause);
        }
    }
}
