from __future__ import annotations

from llama_index.core.embeddings import MockEmbedding

from app.modules.data_process.document_chunking import (
    DocumentChunk,
    _compact_with_offsets,
    _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_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)