scripts/verify_chunker.py

"""Comprehensive verification test suite for multimodal PDF element chunker.

Tests compliance with the CocoIndex PDF Elements architecture and LlamaIndex
RecursiveTextSplitter integration:

1. Text splitting using LlamaIndex SentenceSplitter with boundary preservation.
2. Table chunking with header repetition on large partitioned financial tables.
3. Figure/chart chunking with pixel filtering, classification, and image path tracking.
4. Heading breadcrumb propagation and page coordinates.
5. Deterministic fingerprinting and JSON serialization round-trips.
6. Real filing extraction processing (Annual Report, Presentation, Transcript).

Google Python Style Guide Compliant.
"""

from __future__ import annotations

import json
from pathlib import Path
import sys
import unittest

from ingestion.chunker import (
    CHUNKER_VERSION,
    DEFAULT_ANNUAL_SKIP_SECTIONS,
    ELEMENT_FIGURE,
    ELEMENT_TABLE,
    ELEMENT_TEXT,
    Chunk,
    ChunkSet,
    chunk_document,
    read_chunk_cache,
    write_chunk_cache,
)
from ingestion.documents.content import (
    KIND_FIGURE,
    KIND_HEADING,
    KIND_LIST,
    KIND_TABLE,
    KIND_TEXT,
    Block,
    ExtractedDocument,
    Figure,
    Table,
)


class TestSentenceSplitter(unittest.TestCase):
    """Verifies recursive text splitting logic."""

    def test_small_text_stays_single_chunk(self) -> None:
        doc = ExtractedDocument(
            doc_id="test_doc",
            doc_type="concall_presentation",
            ticker="TEST",
            label="Test Presentation",
            blocks=[
                Block(kind=KIND_HEADING, text="Overview", path=["Overview"], page=1),
                Block(kind=KIND_TEXT, text="Short business update paragraph.", path=["Overview"], page=1),
            ],
        )
        chunk_set = chunk_document(doc, chunk_size=500, chunk_overlap=50)
        self.assertEqual(chunk_set.n_chunks, 1)
        chunk = chunk_set.chunks[0]
        self.assertEqual(chunk.element_type, ELEMENT_TEXT)
        self.assertEqual(chunk.page_start, 1)
        self.assertEqual(chunk.path, ["Overview"])
        self.assertIn("Short business update paragraph.", chunk.text)
        self.assertIn("Section: Overview", chunk.embed_text)

    def test_long_prose_is_recursively_split(self) -> None:
        long_paragraph = (
            "Operating margins improved across all three business units during the fiscal period. "
            * 25
        )
        doc = ExtractedDocument(
            doc_id="test_long",
            doc_type="annual_report",
            ticker="TEST",
            label="Annual Report",
            blocks=[
                Block(kind=KIND_HEADING, text="MD&A", path=["MD&A"], page=5),
                Block(kind=KIND_TEXT, text=long_paragraph, path=["MD&A"], page=5),
            ],
        )
        chunk_set = chunk_document(doc, chunk_size=300, chunk_overlap=50)
        self.assertGreater(chunk_set.n_chunks, 1)
        for chunk in chunk_set.chunks:
            self.assertEqual(chunk.element_type, ELEMENT_TEXT)
            self.assertEqual(chunk.page_start, 5)
            self.assertGreater(len(chunk.text), 0)
            self.assertLess(len(chunk.text), 2000)


class TestTableChunking(unittest.TestCase):
    """Verifies structured table preservation and large table partitioning."""

    def test_small_table_stays_intact(self) -> None:
        table = Table(
            rows=[
                ["Segment", "Revenue (INR Cr)", "EBITDA Margin"],
                ["IT Services", "22,500", "16.4%"],
                ["Consulting", "4,200", "14.1%"],
            ],
            header_rows=1,
            caption="Segment Performance",
            page=12,
        )
        doc = ExtractedDocument(
            doc_id="test_tbl_small",
            doc_type="annual_report",
            ticker="TEST",
            label="Annual Report",
            blocks=[
                Block(kind=KIND_HEADING, text="Segment Review", path=["Segment Review"], page=12),
                Block(kind=KIND_TABLE, text=table.to_markdown(), path=["Segment Review"], table=table, page=12),
            ],
        )
        chunk_set = chunk_document(doc)
        self.assertEqual(chunk_set.n_chunks, 1)
        chunk = chunk_set.chunks[0]
        self.assertEqual(chunk.element_type, ELEMENT_TABLE)
        self.assertEqual(chunk.page_start, 12)
        self.assertIn("IT Services", chunk.text)
        self.assertIn("Consulting", chunk.text)
        self.assertEqual(chunk.metadata["n_rows"], 3)
        self.assertEqual(chunk.metadata["header_rows"], 1)

    def test_large_table_is_partitioned_with_headers_repeated(self) -> None:
        header = [["Line Item", "FY26", "FY25"]]
        body = [[f"Schedule Item {i}", f"{1000 + i}", f"{900 + i}"] for i in range(1, 60)]
        table = Table(
            rows=header + body,
            header_rows=1,
            caption="Consolidated Financial Schedules",
            page=20,
        )
        doc = ExtractedDocument(
            doc_id="test_tbl_large",
            doc_type="annual_report",
            ticker="TEST",
            label="Annual Report",
            blocks=[
                Block(kind=KIND_HEADING, text="Financials", path=["Financials"], page=20),
                Block(kind=KIND_TABLE, text=table.to_markdown(), path=["Financials"], table=table, page=20),
            ],
        )
        chunk_set = chunk_document(doc)
        self.assertGreater(chunk_set.n_chunks, 1)
        for chunk in chunk_set.chunks:
            self.assertEqual(chunk.element_type, ELEMENT_TABLE)
            self.assertIn("Line Item", chunk.text)
            self.assertIn("FY26", chunk.text)
            self.assertIn("Table Part", chunk.embed_text)


class TestFigureChunking(unittest.TestCase):
    """Verifies visual element and figure metadata extraction."""

    def test_meaningful_chart_creates_figure_chunk(self) -> None:
        fig = Figure(
            path="figures/test_fig/p0004-01.png",
            page=4,
            kind="bar_chart",
            width=800,
            height=400,
            caption="Quarterly Revenue Trend",
        )
        doc = ExtractedDocument(
            doc_id="test_fig",
            doc_type="concall_presentation",
            ticker="TEST",
            label="Investor Presentation",
            blocks=[
                Block(kind=KIND_HEADING, text="Revenue Growth", path=["Revenue Growth"], page=4),
                Block(kind=KIND_FIGURE, text="", path=["Revenue Growth"], figure=fig, page=4),
            ],
        )
        chunk_set = chunk_document(doc)
        self.assertEqual(chunk_set.n_chunks, 1)
        chunk = chunk_set.chunks[0]
        self.assertEqual(chunk.element_type, ELEMENT_FIGURE)
        self.assertEqual(chunk.page_start, 4)
        self.assertEqual(chunk.figure_path, "figures/test_fig/p0004-01.png")
        self.assertIn("figures/test_fig/p0004-01.png", chunk.text)
        self.assertEqual(chunk.metadata["kind"], "bar_chart")
        self.assertEqual(chunk.metadata["width"], 800)
        self.assertEqual(chunk.metadata["height"], 400)

    def test_tiny_decorative_furniture_is_filtered(self) -> None:
        tiny_icon = Figure(
            path="figures/test_fig/icon.png",
            page=1,
            kind="icon",
            width=24,
            height=24,
            caption="",
        )
        doc = ExtractedDocument(
            doc_id="test_filter",
            doc_type="concall_presentation",
            ticker="TEST",
            label="Presentation",
            blocks=[
                Block(kind=KIND_FIGURE, text="", path=[], figure=tiny_icon, page=1),
            ],
        )
        chunk_set = chunk_document(doc, min_figure_pixels=48)
        self.assertEqual(chunk_set.n_chunks, 0)


class TestChunkSerialization(unittest.TestCase):
    """Verifies chunk serialization and persistence."""

    def test_chunk_and_chunkset_roundtrip(self) -> None:
        chunk = Chunk(
            chunk_id="test_p0001_txt_001",
            doc_id="test_doc",
            ticker="TEST",
            doc_type="annual_report",
            label="FY26 Report",
            element_type=ELEMENT_TEXT,
            page_start=1,
            page_end=1,
            section="Executive Summary",
            path=["Executive Summary", "Strategic Highlights"],
            text="Revenue increased 12% YoY.",
            embed_text="[TEST] Revenue increased 12% YoY.",
            metadata={"char_count": 26},
        )
        chunk_set = ChunkSet(
            doc_id="test_doc",
            ticker="TEST",
            doc_type="annual_report",
            label="FY26 Report",
            fingerprint="abc123hash",
            params={"chunk_size": 800},
            chunks=[chunk],
        )

        serialized = chunk_set.to_dict()
        reconstructed = ChunkSet.from_dict(serialized)

        self.assertEqual(reconstructed.doc_id, chunk_set.doc_id)
        self.assertEqual(reconstructed.fingerprint, chunk_set.fingerprint)
        self.assertEqual(reconstructed.n_chunks, 1)
        r_chunk = reconstructed.chunks[0]
        self.assertEqual(r_chunk.chunk_id, chunk.chunk_id)
        self.assertEqual(r_chunk.element_type, chunk.element_type)
        self.assertEqual(r_chunk.path, chunk.path)

    def test_llamaindex_node_conversion(self) -> None:
        chunk_txt = Chunk(
            chunk_id="test_txt_001",
            doc_id="test_doc",
            ticker="TEST",
            element_type=ELEMENT_TEXT,
            page_start=2,
            page_end=2,
            section="Overview",
            text="Prose content for LlamaIndex node.",
        )
        chunk_fig = Chunk(
            chunk_id="test_fig_001",
            doc_id="test_doc",
            ticker="TEST",
            element_type=ELEMENT_FIGURE,
            page_start=4,
            page_end=4,
            figure_path="figures/test_doc/p0004-01.png",
            text="Figure markdown description.",
        )
        chunk_set = ChunkSet(
            doc_id="test_doc",
            ticker="TEST",
            chunks=[chunk_txt, chunk_fig],
        )

        nodes = chunk_set.to_nodes()
        self.assertEqual(len(nodes), 2)
        from llama_index.core.schema import ImageNode, TextNode
        self.assertIsInstance(nodes[0], TextNode)
        self.assertIsInstance(nodes[1], ImageNode)
        self.assertEqual(nodes[1].image_path, "figures/test_doc/p0004-01.png")

        reconstructed_chunk = Chunk.from_node(nodes[0])
        self.assertEqual(reconstructed_chunk.chunk_id, "test_txt_001")
        self.assertEqual(reconstructed_chunk.text, "Prose content for LlamaIndex node.")

        from llama_index.core.schema import NodeRelationship
        self.assertIn(NodeRelationship.NEXT, nodes[0].relationships)
        self.assertEqual(nodes[0].relationships[NodeRelationship.NEXT].node_id, nodes[1].node_id)
        self.assertIn(NodeRelationship.PREVIOUS, nodes[1].relationships)
        self.assertEqual(nodes[1].relationships[NodeRelationship.PREVIOUS].node_id, nodes[0].node_id)

        reconstructed_set = ChunkSet.from_nodes(nodes, doc_id="test_doc", ticker="TEST")
        self.assertEqual(reconstructed_set.n_chunks, 2)


class TestRealExtractions(unittest.TestCase):
    """Verifies chunking against real extracted corporate filings."""

    def setUp(self) -> None:
        self.extracted_dir = Path("output/WIPRO/extracted")

    def test_presentation_chunking(self) -> None:
        path = self.extracted_dir / "presentation_2026_01.json"
        if not path.exists():
            self.skipTest("Extracted presentation not available")
        doc = ExtractedDocument.from_dict(json.loads(path.read_text("utf-8")))
        chunk_set = chunk_document(doc)
        counts = chunk_set.element_counts()
        self.assertGreater(chunk_set.n_chunks, 20)
        self.assertGreater(counts.get(ELEMENT_FIGURE, 0), 10)
        self.assertGreater(counts.get(ELEMENT_TEXT, 0), 10)

    def test_annual_report_chunking(self) -> None:
        path = self.extracted_dir / "annual_report_FY2026_financials.json"
        if not path.exists():
            self.skipTest("Extracted annual report not available")
        doc = ExtractedDocument.from_dict(json.loads(path.read_text("utf-8")))
        chunk_set = chunk_document(doc)
        counts = chunk_set.element_counts()
        self.assertGreater(chunk_set.n_chunks, 30)
        self.assertGreater(counts.get(ELEMENT_TABLE, 0), 5)
        self.assertGreater(counts.get(ELEMENT_TEXT, 0), 30)

    def test_transcript_chunking(self) -> None:
        path = self.extracted_dir / "transcript_2026_07.json"
        if not path.exists():
            self.skipTest("Extracted transcript not available")
        doc = ExtractedDocument.from_dict(json.loads(path.read_text("utf-8")))
        chunk_set = chunk_document(doc)
        self.assertGreater(chunk_set.n_chunks, 20)


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