"""
Semantic Chunker for Word Documents using Docling HybridChunker.

This module provides document chunking capabilities for Word documents using:
- HybridChunker from docling_core for intelligent content segmentation
- Provider-aware tokenizer (tiktoken for Azure, HF AutoTokenizer for vLLM)
- Custom serializers for tables and images
- Contextualization for each chunk

Architecture:
1. DoclingDocument → HybridChunker → Contextualized Chunks
2. Preserves tables in markdown format within chunks
3. Includes image captions and descriptions
4. Token budgeting via the configured embedding provider's tokenizer

Pattern follows: /tmp/pdf_parser_service.py
"""

from typing import Any, Dict, List
from uuid import UUID

import structlog

from docling_core.transforms.chunker.base import BaseChunk
from docling_core.transforms.chunker.hierarchical_chunker import (
    DocChunk,
    ChunkingDocSerializer,
    ChunkingSerializerProvider,
)
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
from docling_core.transforms.serializer.markdown import (
    MarkdownTableSerializer,
    MarkdownPictureSerializer,
)

from src.llm.factory import get_docling_tokenizer
from docling_core.types.doc.document import (
    PictureClassificationData,
    PictureDescriptionData,
    PictureItem,
    DoclingDocument,
)
from docling_core.transforms.serializer.base import (
    BaseDocSerializer,
    SerializationResult,
)
from docling_core.transforms.serializer.common import create_ser_result

logger = structlog.get_logger(__name__)


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.

        Args:
            item: Picture item to serialize
            doc_serializer: Document serializer for post-processing
            doc: Parent Docling document
            **kwargs: Additional serialization arguments

        Returns:
            SerializationResult with formatted text and span source
        """
        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.

        Args:
            doc: Docling document to serialize

        Returns:
            ChunkingDocSerializer configured with custom serializers
        """
        return ChunkingDocSerializer(
            doc=doc,
            picture_serializer=AnnotationPictureSerializer(),
            table_serializer=MarkdownTableSerializer(),
        )


class SemanticChunker:
    """
    Semantic chunker for Word documents using HybridChunker.

    Uses:
    - Provider-aware tokenizer (tiktoken for Azure, HF AutoTokenizer for vLLM)
    - HybridChunker for intelligent document segmentation
    - Custom serializers for tables and images
    - Contextualization for enhanced chunk quality

    Example:
        chunker = SemanticChunker()
        chunks = chunker.chunk_document(docling_doc, file_id, "report.docx")
    """

    def __init__(self, max_tokens: int = 8000):
        """
        Initialize semantic chunker with tokenizer and hybrid chunker.

        Args:
            max_tokens: Maximum tokens per chunk (default: 8000, matching typical
                8K-context embedding models).
        """
        self.logger = structlog.get_logger()

        self.tokenizer = get_docling_tokenizer(max_tokens)

        # Initialize HybridChunker with custom serializers
        self.chunker = HybridChunker(
            tokenizer=self.tokenizer,
            serializer_provider=ImgTableAnnotationSerializerProvider(),
        )

        self.logger.info(
            "SemanticChunker initialized",
            max_tokens=max_tokens,
            tokenizer=type(self.tokenizer).__name__,
        )

    def chunk_document(
        self,
        doc: DoclingDocument,
        file_id: UUID,
        original_filename: str,
    ) -> List[Dict[str, Any]]:
        """
        Chunk a Docling document into semantic units with contextualization.

        Args:
            doc: DoclingDocument from Docling conversion
            file_id: UUID of the file record
            original_filename: Original filename (e.g., "report.docx")

        Returns:
            List of chunk records with text, metadata, and contextualization
        """
        self.logger.info(
            "Starting document chunking",
            filename=original_filename,
            file_id=str(file_id)
        )

        # Apply HybridChunker
        chunk_iter = self.chunker.chunk(dl_doc=doc)
        chunks = list(chunk_iter)

        self.logger.info(f"Created {len(chunks)} raw chunks")

        # Convert chunks to records with contextualization
        chunk_records = []
        for chunk in chunks:
            record = self._build_chunk_record(
                chunk=chunk,
                file_id=file_id,
                original_filename=original_filename
            )
            if record:
                chunk_records.append(record)

        self.logger.info(
            f"Chunking complete",
            total_chunks=len(chunk_records),
            filename=original_filename
        )

        return chunk_records

    def _build_chunk_record(
        self,
        chunk: BaseChunk,
        file_id: UUID,
        original_filename: str,
    ) -> Dict[str, Any] | None:
        """
        Build a chunk record from a Docling chunk with metadata.

        Args:
            chunk: BaseChunk from HybridChunker
            file_id: UUID of the file record
            original_filename: Original filename

        Returns:
            Dictionary with chunk data or None if chunk is empty
        """
        # Get contextualized text
        text = self.chunker.contextualize(chunk=chunk)

        # Skip empty chunks
        if not text.strip():
            return None

        # Validate chunk as DocChunk to access metadata
        doc_chunk = DocChunk.model_validate(chunk)

        # Extract page number from doc items (for multi-sheet/multi-page docs)
        page_numbers = []
        for item in doc_chunk.meta.doc_items:
            if hasattr(item, 'prov') and item.prov:
                for prov_item in item.prov:
                    if hasattr(prov_item, 'page_no') and prov_item.page_no is not None:
                        page_numbers.append(prov_item.page_no)

        # Build metadata
        metadata = {
            'page_number': min(page_numbers) if page_numbers else None,
            'source_filename': original_filename,
            'unique_filename': original_filename,
            'contains_table': self._contains_markdown_table(text),
        }

        # Build chunk record
        record = {
            'file_id': str(file_id),
            'text': text,
            'metadata': metadata,
            'element_type': 'Text',
        }

        return record

    def _contains_markdown_table(self, text: str) -> bool:
        """
        Check if text contains a markdown table.

        Args:
            text: Text content to check

        Returns:
            True if markdown table detected, False otherwise
        """
        # Simple heuristic: look for table separator lines (e.g., "| --- | --- |")
        lines = text.split('\n')
        for line in lines:
            # Check for markdown table separator pattern
            if '|' in line and '---' in line:
                return True
        return False
