from __future__ import annotations

import sys
from types import ModuleType, SimpleNamespace

from llama_index.core.embeddings import MockEmbedding

from app.modules.data_process.document_chunking import (
    DocumentChunk,
    _compact_with_offsets,
    _document_converter,
    _project_layout_span,
    chunk_fixed_text,
    chunk_semantic_text,
    merge_short_chunks,
)


def test_fixed_splitter_preserves_offsets_and_token_limit() -> None:
    text = "第一段说明苹果。第二段说明香蕉。\n第三段说明数据库。第四段说明索引。"
    chunks = chunk_fixed_text(text, chunk_size=20, chunk_overlap=0)

    assert len(chunks) > 1
    assert all(chunk.source_start is not None for chunk in chunks)
    assert all(chunk.source_end is not None for chunk in chunks)
    assert all(
        chunk.original_content == text[chunk.source_start : chunk.source_end]
        for chunk in chunks
        if chunk.source_start is not None and chunk.source_end is not None
    )
    assert all(chunk.token_count <= 20 for chunk in chunks)


def test_semantic_splitter_uses_llamaindex_and_reapplies_maximum_size() -> None:
    text = "第一段讨论水果。第二段继续讨论香蕉。第三段讨论数据库。第四段讨论索引。"
    chunks = chunk_semantic_text(
        text,
        chunk_size=30,
        chunk_overlap=0,
        breakpoint_percentile_threshold=95,
        embed_model=MockEmbedding(embed_dim=8),
    )

    assert len(chunks) >= 2
    assert all(chunk.token_count <= 30 for chunk in chunks)
    assert "".join(chunk.original_content for chunk in chunks) == text


def test_layout_converter_disables_ocr(monkeypatch) -> None:
    class FakePipelineOptions:
        def __init__(self) -> None:
            self.do_ocr = True

    class FakePdfFormatOption:
        def __init__(self, *, pipeline_options) -> None:
            self.pipeline_options = pipeline_options

    class FakeDocumentConverter:
        def __init__(self, *, format_options) -> None:
            self.format_options = format_options

    docling_module = ModuleType("docling")
    docling_module.__path__ = []
    document_converter_module = ModuleType("docling.document_converter")
    document_converter_module.DocumentConverter = FakeDocumentConverter
    document_converter_module.PdfFormatOption = FakePdfFormatOption
    datamodel_module = ModuleType("docling.datamodel")
    datamodel_module.__path__ = []
    base_models_module = ModuleType("docling.datamodel.base_models")
    base_models_module.InputFormat = SimpleNamespace(PDF="pdf")
    pipeline_options_module = ModuleType("docling.datamodel.pipeline_options")
    pipeline_options_module.PdfPipelineOptions = FakePipelineOptions
    for name, module in {
        "docling": docling_module,
        "docling.document_converter": document_converter_module,
        "docling.datamodel": datamodel_module,
        "docling.datamodel.base_models": base_models_module,
        "docling.datamodel.pipeline_options": pipeline_options_module,
    }.items():
        monkeypatch.setitem(sys.modules, name, module)

    _document_converter.cache_clear()
    try:
        converter = _document_converter()
        options = converter.format_options["pdf"].pipeline_options
        assert options.do_ocr is False
    finally:
        _document_converter.cache_clear()


def test_layout_projection_ignores_layout_whitespace_but_keeps_source_lines() -> None:
    source = "标题\n第一条 这是正文。\n第二条 后续正文。"
    compact_source, offsets = _compact_with_offsets(source)
    start, end, cursor = _project_layout_span(
        source,
        "第一条\n这是正文。",
        compact_source=compact_source,
        source_offsets=offsets,
        compact_start=0,
    )

    assert source[start:end] == "第一条 这是正文。"
    assert cursor > 0


def test_short_layout_chunk_merges_with_neighbor_and_keeps_page_provenance() -> None:
    source = "短标题\n这是一段足够长的正文内容，用于测试相邻切片合并。"
    chunks = [
        DocumentChunk("短标题", "短标题", 0, 3, 1, 1, 2, source_pages=(1,)),
        DocumentChunk(
            "这是一段足够长的正文内容，用于测试相邻切片合并。",
            "这是一段足够长的正文内容，用于测试相邻切片合并。",
            4,
            len(source),
            2,
            2,
            20,
            source_pages=(1, 2),
        ),
    ]

    merged = merge_short_chunks(
        chunks,
        source_text=source,
        min_token_count=10,
        max_token_count=100,
    )

    assert len(merged) == 1
    assert merged[0].original_content == source
    assert merged[0].source_pages == (1, 2)