package tech.easyflow.ai.documentimport.task;

import org.springframework.http.MediaType;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.web.context.request.async.AsyncRequestNotUsableException;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tech.easyflow.ai.documentimport.DocumentImportKeys;
import tech.easyflow.ai.entity.Document;
import tech.easyflow.ai.mapper.DocumentMapper;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.annotation.Resource;
import java.io.IOException;
import java.math.BigInteger;
import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 知识库文档任务状态 SSE 推送服务。
 *
 * @author Codex
 * @since 2026-04-15
 */
@Service
public class DocumentImportTaskStatusStreamService {

    private static final Logger LOG = LoggerFactory.getLogger(DocumentImportTaskStatusStreamService.class);
    private static final long SSE_TIMEOUT_MS = Duration.ofMinutes(30).toMillis();

    private final Map<String, Set<SseEmitter>> knowledgeEmitters = new ConcurrentHashMap<String, Set<SseEmitter>>();

    @Resource
    private DocumentMapper documentMapper;

    @Resource(name = "sseThreadPool")
    private ThreadPoolTaskExecutor sseThreadPool;

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    @Resource
    private DocumentImportStatusBroadcastProperties statusBroadcastProperties;

    /**
     * 订阅知识库文档任务状态流。
     *
     * @param knowledgeId 知识库 ID
     * @return SSE 连接
     */
    public SseEmitter subscribe(BigInteger knowledgeId) {
        if (knowledgeId == null) {
            throw new BusinessException("知识库id不能为空");
        }
        String topicKey = toTopicKey(knowledgeId);
        SseEmitter emitter = new SseEmitter(SSE_TIMEOUT_MS);
        knowledgeEmitters.computeIfAbsent(topicKey, key -> ConcurrentHashMap.newKeySet()).add(emitter);
        emitter.onCompletion(() -> removeEmitter(topicKey, emitter));
        emitter.onTimeout(() -> {
            removeEmitter(topicKey, emitter);
            emitter.complete();
        });
        emitter.onError(error -> removeEmitter(topicKey, emitter));
        sendAsync(topicKey, emitter, "connected", buildConnectedPayload(knowledgeId));
        return emitter;
    }

    /**
     * 在事务提交后推送文档任务状态变更。
     *
     * @param documentId 文档 ID
     */
    public void publishAfterCommit(BigInteger documentId) {
        if (documentId == null) {
            return;
        }
        Runnable publishAction = () -> publishStatusChange(documentId);
        if (TransactionSynchronizationManager.isSynchronizationActive()
            && TransactionSynchronizationManager.isActualTransactionActive()) {
            TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {
                @Override
                public void afterCommit() {
                    publishAction.run();
                }
            });
            return;
        }
        publishAction.run();
    }

    /**
     * 处理 Redis 广播收到的文档状态变更。
     *
     * @param documentId 文档 ID
     */
    public void publishLocal(BigInteger documentId) {
        publishNow(documentId);
    }

    private void publishStatusChange(BigInteger documentId) {
        // 先推送本机连接，降低单机部署和广播链路延迟。
        publishNow(documentId);
        stringRedisTemplate.convertAndSend(statusBroadcastProperties.getStatusBroadcastChannel(), documentId.toString());
    }

    void publishNow(BigInteger documentId) {
        Document document = documentMapper.selectOneById(documentId);
        if (document == null || document.getCollectionId() == null) {
            return;
        }
        String topicKey = toTopicKey(document.getCollectionId());
        Set<SseEmitter> emitters = knowledgeEmitters.get(topicKey);
        if (emitters == null || emitters.isEmpty()) {
            return;
        }
        Map<String, Object> payload = buildDocumentPayload(document);
        for (SseEmitter emitter : emitters) {
            sendAsync(topicKey, emitter, "document-status", payload);
        }
    }

    private Map<String, Object> buildConnectedPayload(BigInteger knowledgeId) {
        Map<String, Object> payload = new LinkedHashMap<String, Object>();
        payload.put("knowledgeId", knowledgeId.toString());
        payload.put("type", "connected");
        return payload;
    }

    private Map<String, Object> buildDocumentPayload(Document document) {
        Map<String, Object> payload = new LinkedHashMap<String, Object>();
        payload.put("type", "document-status");
        payload.put("knowledgeId", document.getCollectionId() == null ? null : document.getCollectionId().toString());
        payload.put("documentId", document.getId() == null ? null : document.getId().toString());
        payload.put("processStatus", document.getProcessStatus());
        payload.put("progressPercent", document.getProgressPercent());
        payload.put("totalChunks", document.getTotalChunks());
        payload.put("completedChunks", document.getCompletedChunks());
        payload.put("failedChunks", document.getFailedChunks());
        payload.put("parseCurrentStage", readOptionAsString(document, DocumentImportKeys.KEY_DOCUMENT_PARSE_CURRENT_STAGE));
        payload.put("parseStatusMessage", readOptionAsString(document, DocumentImportKeys.KEY_DOCUMENT_PARSE_STATUS_MESSAGE));
        payload.put("lastTaskError", document.getLastTaskError());
        payload.put("taskModifiedAt", document.getTaskModifiedAt());
        return payload;
    }

    private String readOptionAsString(Document document, String key) {
        if (document == null || document.getOptions() == null || key == null) {
            return null;
        }
        Object value = document.getOptions().get(key);
        return value == null ? null : String.valueOf(value);
    }

    private void sendAsync(String topicKey, SseEmitter emitter, String eventName, Map<String, Object> payload) {
        sseThreadPool.execute(() -> {
            if (!isEmitterRegistered(topicKey, emitter)) {
                return;
            }
            try {
                emitter.send(
                    SseEmitter.event()
                        .name(eventName)
                        .data(payload, MediaType.APPLICATION_JSON)
                );
            } catch (Exception e) {
                removeEmitter(topicKey, emitter);
                if (isClientDisconnected(e)) {
                    LOG.debug("文档导入状态流客户端已断开: topicKey={}, eventName={}, message={}",
                        topicKey, eventName, e.getMessage());
                    return;
                }
                LOG.warn("文档导入状态流推送失败: topicKey={}, eventName={}", topicKey, eventName, e);
                completeQuietly(emitter);
            }
        });
    }

    /**
     * 判断指定 SSE 连接是否仍注册在主题下，避免已清理连接继续被异步任务写入。
     *
     * @param topicKey 主题键
     * @param emitter SSE 连接
     * @return 是否仍处于注册状态
     */
    private boolean isEmitterRegistered(String topicKey, SseEmitter emitter) {
        Set<SseEmitter> emitters = knowledgeEmitters.get(topicKey);
        return emitters != null && emitters.contains(emitter);
    }

    private void removeEmitter(String topicKey, SseEmitter emitter) {
        Set<SseEmitter> emitters = knowledgeEmitters.get(topicKey);
        if (emitters == null) {
            return;
        }
        emitters.remove(emitter);
        if (emitters.isEmpty()) {
            knowledgeEmitters.remove(topicKey);
        }
    }

    /**
     * 判断异常是否由客户端断开 SSE 连接导致。
     *
     * @param throwable 异常
     * @return 是否为客户端断连
     */
    private boolean isClientDisconnected(Throwable throwable) {
        Throwable current = throwable;
        while (current != null) {
            if (current instanceof AsyncRequestNotUsableException || current instanceof IOException) {
                return true;
            }
            String message = current.getMessage();
            if (message != null) {
                String lowerMessage = message.toLowerCase();
                if (lowerMessage.contains("broken pipe")
                    || lowerMessage.contains("connection reset")
                    || lowerMessage.contains("response not usable")
                    || lowerMessage.contains("client abort")) {
                    return true;
                }
            }
            current = current.getCause();
        }
        return false;
    }

    /**
     * 安静关闭 SSE 连接。
     *
     * @param emitter SSE 连接
     */
    private void completeQuietly(SseEmitter emitter) {
        try {
            emitter.complete();
        } catch (Exception e) {
            LOG.debug("关闭文档导入状态流失败: message={}", e.getMessage());
        }
    }

    private String toTopicKey(BigInteger knowledgeId) {
        return String.valueOf(knowledgeId);
    }
}
