"""Unit tests for the section-aware chunker."""

from __future__ import annotations

from extractor.chunker import chunk_elements
from extractor.schema import SCHEMA_VERSION, BaseElement, element_id


def make_element(
    text: str,
    element_type: str = "text:narrative",
    section_path: list[str] | None = None,
    heading_level: int | None = None,
    seq: int = 0,
) -> BaseElement:
    eid = element_id(text, None, seq, "test_sha256")
    return BaseElement(
        id=eid,
        element_type=element_type,
        text=text,
        section_path=section_path or [],
        section_path_tier=1,
        sequence_index=seq,
        extraction_method="test",
        schema_version=SCHEMA_VERSION,
        source_filename="test.txt",
        source_sha256="test_sha256",
        source_media_type="text/plain",
    )


SOURCE = dict(
    source_filename="test.txt",
    source_sha256="test_sha256",
    source_media_type="text/plain",
)


class TestChunkerBasic:
    def test_empty_input_produces_no_chunks(self):
        chunks = list(chunk_elements(iter([]), **SOURCE))
        assert chunks == []

    def test_single_element_becomes_chunk(self):
        el = make_element("Hello world.", seq=0)
        chunks = list(chunk_elements(iter([el]), **SOURCE))
        assert len(chunks) == 1
        assert chunks[0].element_type == "composite:chunk"
        assert "Hello world." in chunks[0].text

    def test_chunk_token_count_ge_1(self):
        els = [make_element("Some text here.", seq=i) for i in range(5)]
        chunks = list(chunk_elements(iter(els), **SOURCE))
        for chunk in chunks:
            assert chunk.token_count >= 1

    def test_chunk_has_required_provenance_fields(self):
        el = make_element("Text.", seq=0)
        chunks = list(chunk_elements(iter([el]), **SOURCE))
        assert len(chunks) == 1
        assert chunks[0].source_filename == "test.txt"
        assert chunks[0].source_sha256 == "test_sha256"
        assert chunks[0].source_media_type == "text/plain"


class TestChunkerSectionBoundaries:
    def test_heading_triggers_flush(self):
        """A new heading should start a new chunk."""
        heading = make_element("Chapter 1", element_type="structural:section_header", seq=0)
        body1 = make_element("Content of chapter 1.", seq=1)
        heading2 = make_element("Chapter 2", element_type="structural:section_header", seq=2)
        body2 = make_element("Content of chapter 2.", seq=3)

        chunks = list(chunk_elements(iter([heading, body1, heading2, body2]), **SOURCE))
        # Should have at least 2 chunks (one per section)
        assert len(chunks) >= 2

    def test_overlap_does_not_cross_section_boundary(self):
        """With overlap enabled, section-boundary flushes must NOT carry overlap."""
        heading1 = make_element("Section A", element_type="structural:section_header", seq=0)
        body1 = make_element("Content A " * 20, seq=1)  # enough tokens to be in overlap
        heading2 = make_element("Section B", element_type="structural:section_header", seq=2)
        body2 = make_element("Content B.", seq=3)

        chunks = list(
            chunk_elements(
                iter([heading1, body1, heading2, body2]),
                **SOURCE,
                overlap_tokens=50,
            )
        )

        # Find chunk containing Section B content
        section_b_chunks = [c for c in chunks if "Content B." in c.text]
        for chunk in section_b_chunks:
            # Must NOT contain "Content A" — no cross-section overlap
            assert "Content A" not in chunk.text, "Overlap crossed section boundary — this is a bug"


class TestChunkerAtomicElements:
    def test_table_is_not_split(self):
        """A table that fits within max_tokens should be in its own chunk or with surrounding text."""
        body = make_element("Before table.", seq=0)
        table = make_element(
            "| A | B |\n|---|---|\n| 1 | 2 |",
            element_type="table:simple",
            seq=1,
        )
        after = make_element("After table.", seq=2)
        chunks = list(chunk_elements(iter([body, table, after]), **SOURCE, max_tokens=512))
        # Table element must appear in exactly one chunk
        table_in_chunks = [c for c in chunks if "| A | B |" in c.text]
        assert len(table_in_chunks) >= 1

    def test_oversized_atomic_emitted_as_solo_chunk(self):
        """An atomic element exceeding max_tokens must be emitted as a solo chunk."""
        # Create a very long table that exceeds 10 tokens
        long_table_text = "| " + " | ".join([f"Column{i}" for i in range(20)]) + " |\n"
        long_table_text += "|" + "---|" * 20 + "\n"
        for row in range(50):
            long_table_text += "| " + " | ".join([f"cell{row}_{i}" for i in range(20)]) + " |\n"

        table = make_element(long_table_text, element_type="table:simple", seq=0)
        body = make_element("Surrounding text.", seq=1)

        chunks = list(chunk_elements(iter([table, body]), **SOURCE, max_tokens=10))
        # Table should be in its own chunk
        table_chunks = [c for c in chunks if "Column0" in c.text]
        assert len(table_chunks) == 1


class TestChunkerOverlap:
    def test_default_overlap_is_zero(self):
        """Default overlap should be 0 (disabled)."""
        els = [make_element(f"Sentence {i}." * 100, seq=i) for i in range(3)]
        chunks = list(chunk_elements(iter(els), **SOURCE, max_tokens=50))
        for chunk in chunks:
            assert chunk.chunk_metadata.overlap_tokens == 0

    def test_overlap_metadata_recorded(self):
        """When overlap is set, it should be recorded in chunk_metadata."""
        els = [make_element(f"Text content here {i}.", seq=i) for i in range(10)]
        chunks = list(chunk_elements(iter(els), **SOURCE, max_tokens=20, overlap_tokens=5))
        for chunk in chunks:
            assert chunk.chunk_metadata.max_tokens == 20
            assert chunk.chunk_metadata.overlap_tokens == 5
