package com.lingyue.graph.service;

import com.lingyue.graph.entity.TextChunk;
import com.lingyue.graph.repository.TextChunkRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.*;

/**
 * 文本分块服务
 * 将长文本分割成适合向量化的小块
 *
 * @author lingyue
 * @since 2026-01-15
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class TextChunkService {

    private final TextChunkRepository textChunkRepository;

    @Value("${rag.chunk.size:500}")
    private int chunkSize;

    @Value("${rag.chunk.overlap:50}")
    private int chunkOverlap;

    /**
     * 对文本进行分块
     *
     * @param documentId    文档ID
     * @param textStorageId 文本存储ID
     * @param text          待分块的文本
     * @return 分块列表
     */
    @Transactional
    public List<TextChunk> chunkText(String documentId, String textStorageId, String text) {
        if (text == null || text.trim().isEmpty()) {
            log.warn("文本为空，跳过分块: documentId={}", documentId);
            return Collections.emptyList();
        }
        log.info("开始文本分块: documentId={}, textLength={}", documentId, text.length());

        int safeChunkOverlap = chunkOverlap;
        if (safeChunkOverlap < 0) {
            log.warn("分块重叠小于 0，已修正为 0: overlap={}", chunkOverlap);
            safeChunkOverlap = 0;
        }
        if (safeChunkOverlap >= chunkSize) {
            log.warn("分块重叠过大，已修正为 size-1: overlap={}, size={}", chunkOverlap, chunkSize);
            safeChunkOverlap = Math.max(0, chunkSize - 1);
        }

        // 先删除已有的分块
        textChunkRepository.deleteByDocumentId(documentId);

        List<TextChunk> chunks = new ArrayList<>();
        int start = 0;
        int chunkIndex = 0;
        int lastReportedChunk = 0;
        int prevStart = -1;
        int maxChunks = Math.max(1, (text.length() / Math.max(1, (chunkSize - safeChunkOverlap))) + 2);

        while (start < text.length()) {
            if (chunks.size() > maxChunks) {
                log.warn("分块数量异常，提前终止: documentId={}, chunks={}, textLength={}, chunkSize={}, overlap={}",
                        documentId, chunks.size(), text.length(), chunkSize, chunkOverlap);
                break;
            }

            int end = Math.min(start + chunkSize, text.length());

            // 尝试在句子边界分割（避免截断句子）
            if (end < text.length()) {
                int sentenceEnd = findSentenceBoundary(text, start, end);
                if (sentenceEnd > start) {
                    end = sentenceEnd;
                }
            }

            String chunkContent = text.substring(start, end).trim();

            // 跳过空白分块
            if (chunkContent.isEmpty()) {
                start = end;
                continue;
            }

            // 创建分块
            TextChunk chunk = new TextChunk();
            chunk.setId(UUID.randomUUID().toString().replace("-", ""));
            chunk.setDocumentId(documentId);
            chunk.setTextStorageId(textStorageId);
            chunk.setChunkIndex(chunkIndex++);
            chunk.setContent(chunkContent);
            chunk.setTokenCount(estimateTokenCount(chunkContent));

            // 添加元数据
            Map<String, Object> metadata = new HashMap<>();
            metadata.put("start_pos", start);
            metadata.put("end_pos", end);
            metadata.put("length", chunkContent.length());
            chunk.setMetadata(metadata);

            chunk.setCreateTime(new Date());
            chunk.setUpdateTime(new Date());

            textChunkRepository.insert(chunk);
            chunks.add(chunk);

            if (end >= text.length()) {
                break;
            }

            // 移动起始位置（考虑重叠）
            start = end - safeChunkOverlap;
            if (start <= prevStart) {
                log.warn("分块起点未推进，修正并终止: documentId={}, prevStart={}, start={}, end={}, overlap={}",
                        documentId, prevStart, start, end, chunkOverlap);
                start = end;
                break;
            }
            prevStart = start;
            if (start >= text.length() || start <= 0) {
                break;
            }

            // 进度日志：每 20 个分块输出一次
            if (chunks.size() - lastReportedChunk >= 20) {
                log.info("文本分块进度: documentId={}, chunks={}, pos={}/{}",
                        documentId, chunks.size(), end, text.length());
                lastReportedChunk = chunks.size();
            }
        }

        log.info("文档 {} 分块完成，共 {} 块", documentId, chunks.size());
        return chunks;
    }

    /**
     * 寻找句子边界
     * 从end位置向前搜索，找到最近的句子结束符
     */
    private int findSentenceBoundary(String text, int start, int end) {
        // 向前搜索最多100个字符
        int searchStart = Math.max(start, end - 100);

        for (int i = end; i > searchStart; i--) {
            char c = text.charAt(i - 1);
            // 中文和英文的句子结束符
            if (c == '。' || c == '！' || c == '？' || c == '；' ||
                    c == '.' || c == '!' || c == '?' || c == '\n') {
                return i;
            }
        }

        // 如果没找到句子边界，尝试在逗号或空格处分割
        for (int i = end; i > searchStart; i--) {
            char c = text.charAt(i - 1);
            if (c == '，' || c == ',' || c == ' ' || c == '\t') {
                return i;
            }
        }

        return end;
    }

    /**
     * 估算 Token 数量
     * 中文约 1.5-2 字符/token，英文约 4 字符/token
     * 这里使用混合估算
     */
    private int estimateTokenCount(String text) {
        int chineseCount = 0;
        int otherCount = 0;

        for (char c : text.toCharArray()) {
            if (Character.UnicodeScript.of(c) == Character.UnicodeScript.HAN) {
                chineseCount++;
            } else {
                otherCount++;
            }
        }

        // 中文按 1.5 字符/token，其他按 4 字符/token
        return (int) (chineseCount / 1.5 + otherCount / 4.0);
    }

    /**
     * 根据文档ID获取所有分块
     */
    public List<TextChunk> getChunksByDocumentId(String documentId) {
        return textChunkRepository.findByDocumentId(documentId);
    }

    /**
     * 根据文档ID删除所有分块
     */
    @Transactional
    public int deleteByDocumentId(String documentId) {
        int count = textChunkRepository.deleteByDocumentId(documentId);
        log.info("删除文档 {} 的分块，共 {} 块", documentId, count);
        return count;
    }

    /**
     * 统计文档的分块数量
     */
    public int countByDocumentId(String documentId) {
        return textChunkRepository.countByDocumentId(documentId);
    }
}
