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

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

/**
 * 语义化文本切片器
 * 优先在段落边界处切分，其次在句子边界处切分，保持语义完整性
 */
@Slf4j
public class SemanticTextSplitter extends TextSplitter {

    private final int chunkSize;
    private final int chunkOverlap;
    private static final List<String> PARAGRAPH_SEPARATORS = Arrays.asList("\n\n\n", "\n\n", "\n");
    private static final Pattern SENTENCE_END_PATTERN = Pattern.compile("[。！？.!?]+[\\s\"'）)】\\]]*");
    private final MarkdownQaSplitter.TokenEstimator tokenEstimator;

    public SemanticTextSplitter(int chunkSize, int chunkOverlap) {
        this.chunkSize = chunkSize;
        this.chunkOverlap = Math.min(chunkOverlap, chunkSize / 2);
        this.tokenEstimator = new SimpleTokenEstimator();
    }

    public SemanticTextSplitter(int chunkSize) { this(chunkSize, 50); }

    @Override
    protected List<String> splitText(String text) {
        if (StrUtil.isEmpty(text)) return Collections.emptyList();
        return splitTextRecursive(text);
    }

    private List<String> splitTextRecursive(String text) {
        int textTokens = tokenEstimator.estimate(text);
        if (textTokens <= chunkSize) return Collections.singletonList(text.trim());

        List<String> splits = null;
        String usedSeparator = null;
        for (String separator : PARAGRAPH_SEPARATORS) {
            if (text.contains(separator)) { splits = Arrays.asList(text.split(Pattern.quote(separator))); usedSeparator = separator; break; }
        }
        if (splits == null || splits.size() == 1) { splits = splitBySentences(text); usedSeparator = ""; }
        return mergeSplits(splits, usedSeparator);
    }

    private List<String> splitBySentences(String text) {
        List<String> sentences = new ArrayList<>();
        int lastEnd = 0;
        Matcher matcher = SENTENCE_END_PATTERN.matcher(text);
        while (matcher.find()) {
            String sentence = text.substring(lastEnd, matcher.end()).trim();
            if (StrUtil.isNotEmpty(sentence)) sentences.add(sentence);
            lastEnd = matcher.end();
        }
        if (lastEnd < text.length()) { String remaining = text.substring(lastEnd).trim(); if (StrUtil.isNotEmpty(remaining)) sentences.add(remaining); }
        return sentences.isEmpty() ? Collections.singletonList(text) : sentences;
    }

    private List<String> mergeSplits(List<String> splits, String separator) {
        List<String> chunks = new ArrayList<>();
        List<String> currentChunks = new ArrayList<>();
        int currentLength = 0;
        for (String split : splits) {
            if (StrUtil.isEmpty(split)) continue;
            int splitTokens = tokenEstimator.estimate(split);
            if (splitTokens > chunkSize) {
                if (!currentChunks.isEmpty()) { chunks.add(String.join(separator, currentChunks).trim()); currentChunks.clear(); currentLength = 0; }
                if (!separator.isEmpty()) { chunks.addAll(splitTextRecursive(split)); } else { chunks.addAll(forceSplitLongText(split)); }
                continue;
            }
            int separatorTokens = StrUtil.isEmpty(separator) ? 0 : tokenEstimator.estimate(separator);
            if (!currentChunks.isEmpty() && currentLength + splitTokens + separatorTokens > chunkSize) {
                chunks.add(String.join(separator, currentChunks).trim());
                currentChunks = getOverlappingChunks(currentChunks, separator);
                currentLength = estimateTokens(currentChunks, separator);
            }
            currentChunks.add(split);
            currentLength += splitTokens + separatorTokens;
        }
        if (!currentChunks.isEmpty()) chunks.add(String.join(separator, currentChunks).trim());
        return chunks;
    }

    private List<String> getOverlappingChunks(List<String> chunks, String separator) {
        if (chunkOverlap == 0 || chunks.isEmpty()) return new ArrayList<>();
        List<String> overlapping = new ArrayList<>();
        int tokens = 0;
        for (int i = chunks.size() - 1; i >= 0; i--) {
            String chunk = chunks.get(i);
            int chunkTokens = tokenEstimator.estimate(chunk);
            if (tokens + chunkTokens > chunkOverlap) break;
            overlapping.add(0, chunk);
            tokens += chunkTokens + (StrUtil.isEmpty(separator) ? 0 : tokenEstimator.estimate(separator));
        }
        return overlapping;
    }

    private int estimateTokens(List<String> chunks, String separator) {
        int total = 0;
        for (int i = 0; i < chunks.size(); i++) { total += tokenEstimator.estimate(chunks.get(i)); if (i < chunks.size() - 1 && StrUtil.isNotEmpty(separator)) total += tokenEstimator.estimate(separator); }
        return total;
    }

    private List<String> forceSplitLongText(String text) {
        List<String> chunks = new ArrayList<>();
        int charsPerChunk = (int)(chunkSize * 0.8);
        for (int i = 0; i < text.length(); i += charsPerChunk) { int end = Math.min(i + charsPerChunk, text.length()); chunks.add(text.substring(i, end).trim()); }
        return chunks;
    }

    private static class SimpleTokenEstimator implements MarkdownQaSplitter.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);
        }
    }
}