#!/usr/bin/env python3
"""Test HybridChunker with and without custom serializer."""

import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).parent.parent))

from docling.document_converter import DocumentConverter
from docling_core.transforms.chunker import HybridChunker
from docling_core.transforms.chunker.tokenizer import OpenAITokenizer
from docling_core.types import DoclingDocument, PictureItem
from docling_core.types.doc.labels import PictureClassificationData, PictureDescriptionData
from docling_core.transforms.chunker.base_chunker import ChunkingSerializerProvider, create_ser_result, SerializationResult
from docling_core.transforms.chunker.serializer import ChunkingDocSerializer, MarkdownTableSerializer, MarkdownPictureSerializer, BaseDocSerializer
import tiktoken


class AnnotationPictureSerializer(MarkdownPictureSerializer):
    """Custom serializer for picture annotations that includes captions and descriptions."""

    def serialize(
        self,
        *,
        item: PictureItem,
        doc_serializer: BaseDocSerializer,
        doc: DoclingDocument,
        **kwargs: Any,
    ) -> SerializationResult:
        """Serialize picture item with captions and annotations."""
        text_parts: list[str] = []

        # Add caption if available
        caption = item.caption_text(doc=doc)
        if caption:
            text_parts.append(f"[CAPTION]: {caption}")

        # Add annotations
        for annotation in item.annotations:
            if isinstance(annotation, PictureClassificationData):
                predicted_class = (
                    annotation.predicted_classes[0].class_name
                    if annotation.predicted_classes
                    else None
                )
                if predicted_class is not None:
                    text_parts.append(f"[PICTURE TYPE]: {predicted_class}")
            elif isinstance(annotation, PictureDescriptionData):
                text_parts.append(
                    f"[PICTURE DESCRIPTION] {annotation.text} [END OF PICTURE DESCRIPTION]"
                )

        text_res = "\n".join(text_parts)
        text_res = doc_serializer.post_process(text=text_res)
        return create_ser_result(text=text_res, span_source=item)


class ImgTableAnnotationSerializerProvider(ChunkingSerializerProvider):
    """Custom serializer provider for chunking with image and table support."""

    def get_serializer(self, doc: DoclingDocument):
        """Get serializer with custom picture and table serializers."""
        return ChunkingDocSerializer(
            doc=doc,
            picture_serializer=AnnotationPictureSerializer(),
            table_serializer=MarkdownTableSerializer(),
        )


def test_chunker(file_path: str):
    """Test HybridChunker with and without custom serializer."""

    if not Path(file_path).exists():
        print(f"File not found: {file_path}")
        return

    print(f"Testing HybridChunker on: {file_path}\n")

    # Convert with Docling
    converter = DocumentConverter()
    result = converter.convert(Path(file_path))
    doc = result.document

    # Setup tokenizer
    tokenizer = OpenAITokenizer(
        tokenizer=tiktoken.get_encoding("cl100k_base"),
        max_tokens=8000,
    )

    print("="*80)
    print("TEST 1: HybridChunker WITHOUT custom serializer (current)")
    print("="*80)

    chunker1 = HybridChunker(
        tokenizer=tokenizer,
        max_tokens=4000,
        merge_peers=True
    )

    chunks1 = list(chunker1.chunk(doc))
    print(f"\nTotal chunks: {len(chunks1)}")

    # Group by sheet
    sheets1 = {}
    for chunk in chunks1:
        sheet_num = None
        if hasattr(chunk, 'meta') and chunk.meta:
            if hasattr(chunk.meta, 'doc_items') and chunk.meta.doc_items:
                first_item = chunk.meta.doc_items[0]
                if hasattr(first_item, 'prov') and first_item.prov:
                    for prov in first_item.prov:
                        if hasattr(prov, 'page_no'):
                            sheet_num = prov.page_no
                            break

        sheet_key = f"Sheet {sheet_num}" if sheet_num else "Unknown"
        if sheet_key not in sheets1:
            sheets1[sheet_key] = []
        sheets1[sheet_key].append(chunk)

    print("\nChunks by sheet:")
    for sheet, chunks in sorted(sheets1.items()):
        print(f"  {sheet}: {len(chunks)} chunks")

    print("\n" + "="*80)
    print("TEST 2: HybridChunker WITH custom serializer (old version)")
    print("="*80)

    chunker2 = HybridChunker(
        tokenizer=tokenizer,
        max_tokens=4000,
        merge_peers=True,
        serializer_provider=ImgTableAnnotationSerializerProvider(),
    )

    chunks2 = list(chunker2.chunk(doc))
    print(f"\nTotal chunks: {len(chunks2)}")

    # Group by sheet
    sheets2 = {}
    for chunk in chunks2:
        sheet_num = None
        if hasattr(chunk, 'meta') and chunk.meta:
            if hasattr(chunk.meta, 'doc_items') and chunk.meta.doc_items:
                first_item = chunk.meta.doc_items[0]
                if hasattr(first_item, 'prov') and first_item.prov:
                    for prov in first_item.prov:
                        if hasattr(prov, 'page_no'):
                            sheet_num = prov.page_no
                            break

        sheet_key = f"Sheet {sheet_num}" if sheet_num else "Unknown"
        if sheet_key not in sheets2:
            sheets2[sheet_key] = []
        sheets2[sheet_key].append(chunk)

    print("\nChunks by sheet:")
    for sheet, chunks in sorted(sheets2.items()):
        print(f"  {sheet}: {len(chunks)} chunks")

    # Compare
    print("\n" + "="*80)
    print("COMPARISON:")
    print("="*80)
    print(f"WITHOUT custom serializer: {len(chunks1)} chunks")
    print(f"WITH custom serializer: {len(chunks2)} chunks")
    print(f"Difference: {len(chunks2) - len(chunks1)} chunks")

    # Check target sheets
    target_sheets = [2, 3]  # Sheets 2 and 3 (1-based page_no)
    print("\nTarget sheets check:")
    for sheet_num in target_sheets:
        sheet_key = f"Sheet {sheet_num}"
        count1 = len(sheets1.get(sheet_key, []))
        count2 = len(sheets2.get(sheet_key, []))
        print(f"  {sheet_key}:")
        print(f"    Without custom serializer: {count1} chunks")
        print(f"    With custom serializer: {count2} chunks")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python test_hybrid_chunker.py <path_to_xlsx_file>")
        sys.exit(1)

    test_chunker(sys.argv[1])