package com.easyagents.rag.ingestion;

import com.easyagents.rag.core.RagChunk;
import com.easyagents.rag.core.RagChunkTypes;
import com.easyagents.rag.core.RagDefaults;
import com.easyagents.rag.core.RagStrategyCodes;
import com.easyagents.rag.ingestion.analysis.DocumentStructureAnalyzer;
import com.easyagents.rag.ingestion.chunk.RagSplitStrategyRegistry;
import com.easyagents.rag.ingestion.model.AnalysisResult;
import com.easyagents.rag.ingestion.model.StrategyConfig;
import com.easyagents.rag.ingestion.recommend.SplitStrategyRecommender;
import org.junit.Assert;
import org.junit.Test;

import java.util.List;
import java.util.Map;

public class RagIngestionPipelineTest {

    private final DocumentStructureAnalyzer analyzer = new DocumentStructureAnalyzer();
    private final SplitStrategyRecommender recommender = new SplitStrategyRecommender();
    private final RagSplitStrategyRegistry registry = new RagSplitStrategyRegistry();

    @Test
    public void shouldRecommendMarkdownStrategy() {
        String markdown = "# Quick Start\n"
            + "Welcome\n\n"
            + "## Install\n"
            + "Run npm install\n\n"
            + "## Usage\n"
            + "Run pnpm dev";

        AnalysisResult analysis = recommender.recommend(analyzer.analyze(markdown, "md"));

        Assert.assertEquals(RagStrategyCodes.MARKDOWN_SECTION, analysis.getRecommendedStrategyCode());
        Assert.assertTrue(analysis.getConfidence().doubleValue() > 0.4D);
    }

    @Test
    public void shouldRecommendQaStrategyForEnglishAndChinese() {
        String qa = "Q: How to reset password?\n"
            + "A: Open admin page and click reset.\n\n"
            + "问：默认密码是什么？\n"
            + "答：由系统配置统一决定。";

        AnalysisResult analysis = recommender.recommend(analyzer.analyze(qa, "txt"));

        Assert.assertEquals(RagStrategyCodes.QA_PAIR, analysis.getRecommendedStrategyCode());
    }

    @Test
    public void shouldSplitOutlineDocumentByHeadingPath() {
        String outline = "第1章 总则\n适用范围说明。\n\n1.1 目标\n定义系统目标。\n\n1.2 范围\n定义系统范围。";
        AnalysisResult analysis = recommender.recommend(analyzer.analyze(outline, "docx"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.OUTLINE_SECTION);

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertEquals(3, chunks.size());
        Assert.assertEquals("第1章 总则", chunks.get(0).getSourceLabel());
        Assert.assertEquals(2, chunks.get(1).getHeadingPath().size());
        assertHasValidSourceRanges(analysis, chunks.get(0));
    }

    @Test
    public void shouldSplitQaDocumentByPair() {
        String qa = "Q: How to reset password?\n"
            + "A: Open admin page and click reset.\n\n"
            + "问：默认密码是什么？\n"
            + "答：由系统配置统一决定。";
        AnalysisResult analysis = recommender.recommend(analyzer.analyze(qa, "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.QA_PAIR);

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertEquals(2, chunks.size());
        Assert.assertEquals(RagChunkTypes.QA_PAIR, chunks.get(0).getChunkType());
        Assert.assertTrue(chunks.get(0).getContent().contains("问题"));
        Assert.assertTrue(chunks.get(1).getAnswer().contains("系统配置"));
        assertHasValidSourceRanges(analysis, chunks.get(0));
        Assert.assertEquals(2, ((List<?>) chunks.get(0).getOptions().get("sourceRanges")).size());
    }

    @Test
    public void shouldSplitParagraphDocumentWithSourceRanges() {
        String content = "第一段内容用于测试原文映射。\n第二段内容继续补充，便于生成多个分块。\n第三段内容再长一点，确保范围映射稳定。";
        AnalysisResult analysis = recommender.recommend(analyzer.analyze(content, "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH);
        config.setChunkSize(18);
        config.setOverlapSize(4);

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertTrue(chunks.size() > 1);
        assertHasValidSourceRanges(analysis, chunks.get(0));
    }

    @Test
    public void shouldEnforceBgeM3HardLimitForAutoQaTxt() {
        StringBuilder qa = new StringBuilder("问：自动导入为什么失败？\n答：");
        int i = 0;
        while (qa.length() < 77946) {
            qa.append((char) ('一' + i % 20));
            if (i > 0 && i % 700 == 0) {
                qa.append('。');
            }
            i++;
        }
        qa.setLength(77946);
        AnalysisResult analysis = recommender.recommend(analyzer.analyze(qa.toString(), "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.AUTO);

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertEquals(RagStrategyCodes.QA_PAIR, analysis.getRecommendedStrategyCode());
        Assert.assertEquals(77946, analysis.getNormalizedContent().length());
        Assert.assertTrue(chunks.size() > 1);
        assertWithinBgeM3HardLimit(chunks);
        Assert.assertTrue(chunks.stream().allMatch(
            chunk -> chunk.getTokenEstimate().intValue() <= RagDefaults.CHUNK_SIZE));
        Assert.assertTrue(chunks.stream().allMatch(
            chunk -> chunk.getAnswer() != null
                && chunk.getAnswer().length() < analysis.getNormalizedContent().length()));
        assertHasValidSourceRanges(analysis, chunks.get(0));
    }

    @Test
    public void shouldEnforceBgeM3HardLimitForCustomRegexChunk() {
        StringBuilder content = new StringBuilder();
        for (int i = 0; i < 20000; i++) {
            content.append((char) ('甲' + i % 16));
            if (i > 0 && i % 997 == 0) {
                content.append('；');
            }
        }
        AnalysisResult analysis = recommender.recommend(analyzer.analyze(content.toString(), "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.CUSTOM_REGEX);
        config.setRegex("\\|NEVER_MATCH\\|");

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertTrue(chunks.size() > 1);
        assertWithinBgeM3HardLimit(chunks);
        assertHasValidSourceRanges(analysis, chunks.get(0));
    }

    /**
     * Verifies that custom-regex matches remain discarded when the new option is omitted.
     */
    @Test
    public void shouldDiscardCustomRegexMatchByDefault() {
        AnalysisResult analysis = new AnalysisResult();
        analysis.setNormalizedContent(
            "<QUESTION>第一个问题及其答案内容。<QUESTION>第二个问题及其答案内容。");
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.CUSTOM_REGEX);
        config.setRegex("<QUESTION>");

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertEquals(2, chunks.size());
        Assert.assertEquals("第一个问题及其答案内容。", chunks.get(0).getContent());
        Assert.assertEquals("第二个问题及其答案内容。", chunks.get(1).getContent());
    }

    /**
     * Verifies that an enabled option retains each match at the beginning of the next chunk.
     */
    @Test
    public void shouldRetainCustomRegexMatchInNextChunk() {
        AnalysisResult analysis = new AnalysisResult();
        analysis.setNormalizedContent(
            "<QUESTION>第一个问题及其答案内容。<QUESTION>第二个问题及其答案内容。");
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.CUSTOM_REGEX);
        config.setRegex("<QUESTION>");
        config.setRetainRegexMatch(Boolean.TRUE);

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertEquals(2, chunks.size());
        Assert.assertEquals(
            "<QUESTION>第一个问题及其答案内容。",
            chunks.get(0).getContent());
        Assert.assertEquals(
            "<QUESTION>第二个问题及其答案内容。",
            chunks.get(1).getContent());
        assertHasValidSourceRanges(analysis, chunks.get(1));
    }

    @Test
    public void shouldUseChineseAwareTokenEstimate() {
        String content = "知识库自动导入向量化失败，需要检查分块大小和模型上下文窗口。";
        AnalysisResult analysis = recommender.recommend(analyzer.analyze(content, "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH);

        List<RagChunk> chunks = registry.split(analysis, config);

        Assert.assertEquals(Integer.valueOf(content.codePointCount(0, content.length())),
            chunks.get(0).getTokenEstimate());
    }

    @Test
    public void shouldRejectOverlapEqualToChunkSize() {
        AnalysisResult analysis = recommender.recommend(
            analyzer.analyze("知识库分块参数需要保证游标持续向前推进。", "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH);
        config.setChunkSize(128);
        config.setOverlapSize(128);

        try {
            registry.split(analysis, config);
            Assert.fail("overlapSize 等于 chunkSize 时应拒绝分块");
        } catch (IllegalArgumentException expected) {
            Assert.assertTrue(expected.getMessage().contains("overlapSize"));
        }
    }

    @Test
    public void shouldRejectOverlapGreaterThanChunkSize() {
        AnalysisResult analysis = recommender.recommend(
            analyzer.analyze("知识库分块参数需要保证游标持续向前推进。", "txt"));
        StrategyConfig config = StrategyConfig.defaults();
        config.setStrategyCode(RagStrategyCodes.PARAGRAPH_LENGTH);
        config.setChunkSize(128);
        config.setOverlapSize(512);

        try {
            registry.split(analysis, config);
            Assert.fail("overlapSize 大于 chunkSize 时应拒绝分块");
        } catch (IllegalArgumentException expected) {
            Assert.assertTrue(expected.getMessage().contains("overlapSize"));
        }
    }

    private void assertWithinBgeM3HardLimit(List<RagChunk> chunks) {
        for (RagChunk chunk : chunks) {
            Assert.assertNotNull(chunk.getTokenEstimate());
            Assert.assertTrue(
                "chunk token estimate exceeds BGE-M3 hard limit: " + chunk.getTokenEstimate(),
                chunk.getTokenEstimate().intValue() <= RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT);
        }
    }

    @SuppressWarnings("unchecked")
    private void assertHasValidSourceRanges(AnalysisResult analysis, RagChunk chunk) {
        Object rawRanges = chunk.getOptions().get("sourceRanges");
        Assert.assertTrue(rawRanges instanceof List);
        List<Map<String, Object>> ranges = (List<Map<String, Object>>) rawRanges;
        Assert.assertFalse(ranges.isEmpty());
        int normalizedLength = analysis.getNormalizedContent().length();
        for (Map<String, Object> range : ranges) {
            int start = ((Number) range.get("start")).intValue();
            int end = ((Number) range.get("end")).intValue();
            Assert.assertTrue(start >= 0);
            Assert.assertTrue(end > start);
            Assert.assertTrue(end <= normalizedLength);
        }
    }
}
