package com.lingyue.graph.listener;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.lingyue.common.event.DocumentParsedEvent;
import com.lingyue.graph.service.GraphNerService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.event.EventListener;
import org.springframework.http.*;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 文档解析完成事件监听器
 * 监听文档解析完成事件，自动触发 NER 提取并保存到图数据库
 *
 * @author lingyue
 * @since 2026-01-19
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class DocumentParsedEventListener {

    private final GraphNerService graphNerService;
    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;

    @Value("${ner.auto-extract.enabled:true}")
    private boolean nerAutoExtractEnabled;

    @Value("${ner.python-service.url:http://localhost:8001}")
    private String nerServiceUrl;
    
    @Value("${ner.python-service.use-stream:true}")
    private boolean useStreamApi;

    /**
     * 处理文档解析完成事件
     * 异步执行 NER 提取，不阻塞主流程
     */
    @Async
    @EventListener
    public void handleDocumentParsedEvent(DocumentParsedEvent event) {
        if (!nerAutoExtractEnabled) {
            log.debug("NER 自动提取已禁用，跳过: documentId={}", event.getDocumentId());
            return;
        }

        String documentId = event.getDocumentId();
        String userId = event.getUserId();

        log.info("收到文档解析完成事件，开始 NER 提取: documentId={}, userId={}", documentId, userId);

        long startTime = System.currentTimeMillis();

        try {
            // 1. 获取文档文本
            if (!graphNerService.hasDocumentText(documentId)) {
                log.warn("文档文本不存在，跳过 NER: documentId={}", documentId);
                return;
            }
            
            String text = graphNerService.getDocumentText(documentId);
            if (text == null || text.isEmpty()) {
                log.warn("文档文本为空，跳过 NER: documentId={}", documentId);
                return;
            }

            // 2. 调用 Python NER 服务（根据配置选择流式或普通 API）
            Map<String, Object> nerResponse;
            if (useStreamApi) {
                nerResponse = callPythonNerServiceWithStream(documentId, text, userId);
            } else {
                nerResponse = callPythonNerService(documentId, text, userId);
            }
            
            if (nerResponse == null || !Boolean.TRUE.equals(nerResponse.get("success"))) {
                log.warn("NER 服务调用失败: documentId={}, error={}", 
                        documentId, nerResponse != null ? nerResponse.get("errorMessage") : "null response");
                return;
            }

            // 3. 保存实体到图数据库
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> entities = (List<Map<String, Object>>) nerResponse.get("entities");
            Map<String, String> tempIdToNodeId = graphNerService.saveEntitiesToGraph(documentId, userId, entities);

            // 4. 保存关系到图数据库
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> relations = (List<Map<String, Object>>) nerResponse.get("relations");
            int relationCount = graphNerService.saveRelationsToGraph(relations, tempIdToNodeId);

            long processingTime = System.currentTimeMillis() - startTime;
            
            log.info("NER 自动提取完成: documentId={}, entityCount={}, relationCount={}, time={}ms",
                    documentId,
                    entities != null ? entities.size() : 0,
                    relationCount,
                    processingTime);

        } catch (Exception e) {
            log.error("NER 自动提取异常: documentId={}", documentId, e);
            // 异常不向上抛出，不影响其他处理
        }
    }

    /**
     * 调用 Python NER 服务（普通 REST API）
     */
    private Map<String, Object> callPythonNerService(String documentId, String text, String userId) {
        try {
            String url = nerServiceUrl + "/ner/extract";
            
            Map<String, Object> request = new HashMap<>();
            request.put("documentId", documentId);
            request.put("text", text);
            request.put("userId", userId);
            request.put("extractRelations", true);
            
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(request, headers);
            
            ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.POST, entity, Map.class);
            
            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                @SuppressWarnings("unchecked")
                Map<String, Object> body = response.getBody();
                return body;
            }
            
            return null;
            
        } catch (Exception e) {
            log.error("调用 Python NER 服务失败: {}", e.getMessage());
            return null;
        }
    }
    
    /**
     * 调用 Python NER 服务（SSE 流式 API，带进度反馈）
     */
    private Map<String, Object> callPythonNerServiceWithStream(String documentId, String text, String userId) {
        try {
            Map<String, Object> request = new HashMap<>();
            request.put("documentId", documentId);
            request.put("text", text);
            request.put("userId", userId);
            request.put("extractRelations", true);
            
            // 使用 WebClient 处理 SSE
            WebClient webClient = WebClient.builder()
                    .baseUrl(nerServiceUrl)
                    .codecs(configurer -> configurer.defaultCodecs().maxInMemorySize(10 * 1024 * 1024))  // 10MB
                    .build();
            
            AtomicReference<Map<String, Object>> resultRef = new AtomicReference<>();
            
            // 订阅 SSE 事件流，使用 DataBuffer 逐块读取
            StringBuilder buffer = new StringBuilder();
            
            webClient.post()
                    .uri("/ner/extract/stream")
                    .contentType(MediaType.APPLICATION_JSON)
                    .accept(MediaType.TEXT_EVENT_STREAM)
                    .bodyValue(request)
                    .retrieve()
                    .bodyToFlux(String.class)
                    .timeout(Duration.ofMinutes(10))  // 10 分钟总超时
                    .doOnNext(chunk -> {
                        // 累积数据到缓冲区
                        buffer.append(chunk);
                        
                        // 处理完整的 SSE 事件（以双换行符分隔）
                        String bufferStr = buffer.toString();
                        int eventEnd;
                        while ((eventEnd = bufferStr.indexOf("\n\n")) != -1) {
                            String eventBlock = bufferStr.substring(0, eventEnd);
                            bufferStr = bufferStr.substring(eventEnd + 2);
                            buffer.setLength(0);
                            buffer.append(bufferStr);
                            
                            // 解析单个 SSE 事件
                            parseSseEvent(eventBlock, documentId, resultRef);
                        }
                    })
                    .doOnError(e -> log.error("SSE 流处理错误: documentId={}, error={}", documentId, e.getMessage()))
                    .blockLast();  // 阻塞等待完成
            
            // 处理缓冲区中剩余的数据
            if (buffer.length() > 0) {
                parseSseEvent(buffer.toString(), documentId, resultRef);
            }
            
            return resultRef.get();
            
        } catch (Exception e) {
            log.error("调用 Python NER SSE 服务失败: {}", e.getMessage());
            // 回退到普通 API
            log.info("回退到普通 NER API: documentId={}", documentId);
            return callPythonNerService(documentId, text, userId);
        }
    }
    
    /**
     * 解析单个 SSE 事件
     */
    private void parseSseEvent(String eventBlock, String documentId, AtomicReference<Map<String, Object>> resultRef) {
        try {
            String eventType = null;
            String eventData = null;
            
            // 解析事件块
            for (String line : eventBlock.split("\n")) {
                if (line.startsWith("event:")) {
                    eventType = line.substring(6).trim();
                } else if (line.startsWith("data:")) {
                    eventData = line.substring(5).trim();
                }
            }
            
            if (eventData == null || eventData.isEmpty()) {
                return;
            }
            
            @SuppressWarnings("unchecked")
            Map<String, Object> data = objectMapper.readValue(eventData, Map.class);
            
            // 根据事件类型处理
            if ("progress".equals(eventType) || "chunk_complete".equals(eventType)) {
                Object progressObj = data.get("progress_percent");
                if (progressObj != null) {
                    int progress = progressObj instanceof Integer ? (Integer) progressObj : ((Number) progressObj).intValue();
                    String message = (String) data.getOrDefault("message", "处理中...");
                    log.info("NER 进度: documentId={}, progress={}%, message={}", documentId, progress, message);
                }
            } else if ("complete".equals(eventType)) {
                // 完成事件，包含最终结果
                resultRef.set(data);
                log.info("NER 流式处理完成: documentId={}, entities={}", documentId, 
                        data.get("entities") != null ? ((List<?>) data.get("entities")).size() : 0);
            } else if ("error".equals(eventType)) {
                log.error("NER 服务返回错误: documentId={}, error={}", documentId, data.get("error"));
            }
            
        } catch (Exception e) {
            log.debug("解析 SSE 事件时出错: {}", e.getMessage());
        }
    }
}
