package com.easyagents.rag.ingestion;

import com.easyagents.rag.core.RagChunk;
import com.easyagents.rag.core.RagChunkTypes;
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));
    }

    @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);
        }
    }
}
