package cn.iocoder.yudao.module.ai.service.knowledge.splitter;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.transformer.splitter.TextSplitter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Markdown QA 格式专用切片器
 * 识别二级标题（## ）作为问题标记，保持问答对完整性
 */
@Slf4j
public class MarkdownQaSplitter extends TextSplitter {

    private static final Pattern H2_PATTERN = Pattern.compile("^##\\s+(.+)$", Pattern.MULTILINE);
    private static final String PARAGRAPH_SEPARATOR = "\n\n";
    private static final Pattern SENTENCE_PATTERN = Pattern.compile("[。！？.!?]\\s*");

    private final int chunkSize;
    private final TokenEstimator tokenEstimator;

    public MarkdownQaSplitter(int chunkSize) {
        this.chunkSize = chunkSize;
        this.tokenEstimator = new SimpleTokenEstimator();
    }

    @Override
    protected List<String> splitText(String text) {
        if (StrUtil.isEmpty(text)) return Collections.emptyList();
        List<QaPair> qaPairs = parseQaPairs(text);
        if (CollUtil.isEmpty(qaPairs)) return fallbackSplit(text);
        List<String> result = new ArrayList<>();
        for (QaPair qaPair : qaPairs) result.addAll(splitQaPair(qaPair));
        return result;
    }

    private List<QaPair> parseQaPairs(String content) {
        List<QaPair> qaPairs = new ArrayList<>();
        List<Integer> headingPositions = new ArrayList<>();
        List<String> questions = new ArrayList<>();
        Matcher matcher = H2_PATTERN.matcher(content);
        while (matcher.find()) { headingPositions.add(matcher.start()); questions.add(matcher.group(1).trim()); }
        if (CollUtil.isEmpty(headingPositions)) return qaPairs;
        for (int i = 0; i < headingPositions.size(); i++) {
            int start = headingPositions.get(i);
            int end = (i + 1 < headingPositions.size()) ? headingPositions.get(i + 1) : content.length();
            String fullText = content.substring(start, end).trim();
            String question = questions.get(i);
            String answer = fullText.substring(fullText.indexOf('\n') + 1).trim();
            qaPairs.add(new QaPair(question, answer, fullText));
        }
        return qaPairs;
    }

    private List<String> splitQaPair(QaPair qaPair) {
        int qaTokens = tokenEstimator.estimate(qaPair.fullText);
        if (qaTokens <= chunkSize) return Collections.singletonList(qaPair.fullText);
        List<String> chunks = new ArrayList<>();
        List<String> answerChunks = splitLongAnswer(qaPair.answer, qaPair.question);
        for (String answerChunk : answerChunks) {
            chunks.add("## " + qaPair.question + "\n" + answerChunk);
        }
        return chunks;
    }

    private List<String> splitLongAnswer(String answer, String question) {
        List<String> chunks = new ArrayList<>();
        String questionHeader = "## " + question + "\n";
        int questionTokens = tokenEstimator.estimate(questionHeader);
        int availableTokens = chunkSize - questionTokens - 10;
        String[] paragraphs = answer.split(PARAGRAPH_SEPARATOR);
        StringBuilder currentChunk = new StringBuilder();
        int currentTokens = 0;
        for (String paragraph : paragraphs) {
            if (StrUtil.isEmpty(paragraph)) continue;
            int paragraphTokens = tokenEstimator.estimate(paragraph);
            if (paragraphTokens > availableTokens) {
                if (currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
                chunks.addAll(splitLongParagraph(paragraph, availableTokens));
                continue;
            }
            if (currentTokens + paragraphTokens > availableTokens && currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
            if (currentChunk.length() > 0) currentChunk.append("\n\n");
            currentChunk.append(paragraph);
            currentTokens += paragraphTokens;
        }
        if (currentChunk.length() > 0) chunks.add(currentChunk.toString().trim());
        return CollUtil.isEmpty(chunks) ? Collections.singletonList(answer) : chunks;
    }

    private List<String> splitLongParagraph(String paragraph, int availableTokens) {
        List<String> chunks = new ArrayList<>();
        String[] sentences = SENTENCE_PATTERN.split(paragraph);
        StringBuilder currentChunk = new StringBuilder();
        int currentTokens = 0;
        for (String sentence : sentences) {
            if (StrUtil.isEmpty(sentence)) continue;
            int sentenceTokens = tokenEstimator.estimate(sentence);
            if (sentenceTokens > availableTokens) { if (currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; } chunks.add(sentence.trim()); continue; }
            if (currentTokens + sentenceTokens > availableTokens && currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
            currentChunk.append(sentence);
            currentTokens += sentenceTokens;
        }
        if (currentChunk.length() > 0) chunks.add(currentChunk.toString().trim());
        return chunks.isEmpty() ? Collections.singletonList(paragraph) : chunks;
    }

    private List<String> fallbackSplit(String content) {
        List<String> chunks = new ArrayList<>();
        String[] paragraphs = content.split(PARAGRAPH_SEPARATOR);
        StringBuilder currentChunk = new StringBuilder();
        int currentTokens = 0;
        for (String paragraph : paragraphs) {
            if (StrUtil.isEmpty(paragraph)) continue;
            int paragraphTokens = tokenEstimator.estimate(paragraph);
            if (currentTokens + paragraphTokens > chunkSize && currentChunk.length() > 0) { chunks.add(currentChunk.toString().trim()); currentChunk = new StringBuilder(); currentTokens = 0; }
            if (currentChunk.length() > 0) currentChunk.append("\n\n");
            currentChunk.append(paragraph);
            currentTokens += paragraphTokens;
        }
        if (currentChunk.length() > 0) chunks.add(currentChunk.toString().trim());
        return chunks.isEmpty() ? Collections.singletonList(content) : chunks;
    }

    @AllArgsConstructor
    private static class QaPair { String question; String answer; String fullText; }

    public interface TokenEstimator { int estimate(String text); }

    private static class SimpleTokenEstimator implements TokenEstimator {
        @Override
        public int estimate(String text) {
            if (StrUtil.isEmpty(text)) return 0;
            int chineseChars = 0, englishWords = 0;
            for (char c : text.toCharArray()) { if (c >= 0x4E00 && c <= 0x9FA5) chineseChars++; }
            for (String word : text.split("\\s+")) { if (word.matches(".*[a-zA-Z].*")) englishWords++; }
            return chineseChars + (int)(englishWords * 1.3);
        }
    }
}