import unittest

from rag.chunk_config import ChunkConfig
from rag.chunking import chunk_document


class DefaultChunkingTests(unittest.TestCase):
    def test_keeps_short_paragraphs_as_separate_chunks(self):
        result = chunk_document(
            "第一段内容。\n\n第二段内容。",
            "default",
            config=ChunkConfig(size=64, overlap=4),
        )

        self.assertEqual([chunk.text for chunk in result.chunks], ["第一段内容。", "第二段内容。"])
        self.assertEqual(result.actual_strategy, "default")
        self.assertFalse(result.degraded)

    def test_splits_long_paragraphs_with_overlap(self):
        result = chunk_document(
            "abcdefghij" * 3,
            "default",
            config=ChunkConfig(size=10, overlap=2),
        )

        self.assertGreater(len(result.chunks), 1)
        self.assertTrue(all(len(chunk.text) <= 10 for chunk in result.chunks))
        self.assertEqual(result.chunks[0].text[-2:], result.chunks[1].text[:2])


class QaPairChunkingTests(unittest.TestCase):
    def test_keeps_each_question_and_answer_together(self):
        result = chunk_document(
            "Q: 什么是净值？\nA: 净值是基金单位价值。\nQ: 如何申购？\nA: 通过交易页面申购。",
            "qa_pair",
            config=ChunkConfig(size=64, overlap=4),
        )

        self.assertEqual(len(result.chunks), 2)
        self.assertIn("question: 什么是净值？", result.chunks[0].text)
        self.assertIn("answer: 净值是基金单位价值。", result.chunks[0].text)

    def test_rejects_missing_qa_markers(self):
        with self.assertRaises(ValueError):
            chunk_document("普通文本，没有问答标记", "qa_pair")

    def test_rejects_the_whole_document_when_one_pair_is_too_long(self):
        text = "Q: 短问题\nA: 短答案\nQ: 长问题\nA: " + ("很长" * 20)

        with self.assertRaises(ValueError):
            chunk_document(text, "qa_pair", config=ChunkConfig(size=20, overlap=2))


class ChapterChunkingTests(unittest.TestCase):
    def test_preserves_nested_heading_path(self):
        result = chunk_document(
            "# 产品说明\n## 风险揭示\n风险内容。",
            "chapter_semantic",
            config=ChunkConfig(size=64, overlap=4),
        )

        self.assertEqual(result.actual_strategy, "chapter_semantic")
        self.assertEqual(result.chunks[0].section_title, "产品说明 > 风险揭示")
        self.assertTrue(result.chunks[0].text.startswith("【章节：产品说明 > 风险揭示】"))

    def test_splits_only_inside_a_chapter(self):
        result = chunk_document(
            "# 第一章\n" + ("甲" * 20) + "\n# 第二章\n" + ("乙" * 20),
            "chapter_semantic",
            config=ChunkConfig(size=16, overlap=2),
        )

        self.assertTrue(all("甲" not in chunk.text or "乙" not in chunk.text for chunk in result.chunks))

    def test_degrades_to_default_without_headings(self):
        result = chunk_document(
            "没有Markdown标题的普通内容。",
            "chapter_semantic",
            config=ChunkConfig(size=64, overlap=4),
        )

        self.assertEqual(result.actual_strategy, "default")
        self.assertTrue(result.degraded)
        self.assertTrue(result.warning)


if __name__ == "__main__":
    unittest.main()