#!/usr/bin/env python3
"""
Chunking Strategy Evaluation Experiment

This script evaluates three different chunking strategies from docling:
1. Base Chunker (HierarchicalChunker without merging)
2. Hybrid Chunker (Token-aware chunking)
3. Hierarchical Chunker (Structure-based chunking with merging)

The evaluation uses TMUS_992_Q423.pdf to compare how each strategy handles
financial documents with tables.
"""

import sys
from pathlib import Path
from typing import Iterator

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling_core.transforms.chunker.base import BaseChunk
from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker
from docling_core.types.doc import DocItemLabel
from loguru import logger

# Configure logger
logger.remove()
logger.add(sys.stderr, level="INFO")


def setup_converter() -> DocumentConverter:
    """Setup document converter with VLM pipeline for table detection."""
    pipeline_options = PdfPipelineOptions()
    pipeline_options.do_ocr = False
    pipeline_options.do_table_structure = True

    return DocumentConverter(
        format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}
    )


def chunk_with_base(doc) -> list[BaseChunk]:
    """
    Base Chunker: HierarchicalChunker without merging list items.

    This creates one chunk per document element without any merging.
    """
    logger.info("Running Base Chunker (no merging)...")
    chunker = HierarchicalChunker(merge_list_items=False)
    chunks = list(chunker.chunk(dl_doc=doc))
    logger.info(f"Base Chunker produced {len(chunks)} chunks")
    return chunks


def chunk_with_hierarchical(doc) -> list[BaseChunk]:
    """
    Hierarchical Chunker: Respects document structure, merges list items.

    This preserves document hierarchy and combines related elements.
    """
    logger.info("Running Hierarchical Chunker (with merging)...")
    chunker = HierarchicalChunker(merge_list_items=True)
    chunks = list(chunker.chunk(dl_doc=doc))
    logger.info(f"Hierarchical Chunker produced {len(chunks)} chunks")
    return chunks


def chunk_with_hybrid(doc) -> list[BaseChunk]:
    """
    Hybrid Chunker: Token-aware refinement on top of hierarchical chunking.

    This applies tokenization limits while respecting document structure.
    Requires: pip install 'docling-core[chunking]'
    """
    try:
        from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
        from docling_core.transforms.chunker.tokenizer.huggingface import HuggingFaceTokenizer
        from transformers import AutoTokenizer

        logger.info("Running Hybrid Chunker (token-aware)...")

        # Use a common embedding model tokenizer
        # Max tokens set to reasonable size for table context preservation
        tokenizer = HuggingFaceTokenizer(
            tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2"),
            max_tokens=512,  # Suitable for table context
        )

        chunker = HybridChunker(
            tokenizer=tokenizer,
            merge_peers=True,  # Merge undersized chunks with same context
        )

        chunks = list(chunker.chunk(dl_doc=doc))
        logger.info(f"Hybrid Chunker produced {len(chunks)} chunks")
        return chunks

    except ImportError as e:
        logger.error(f"Hybrid chunker dependencies not installed: {e}")
        logger.error("Run: pip install 'docling-core[chunking]' transformers")
        return []


def analyze_chunk(chunk: BaseChunk, chunk_num: int) -> dict:
    """Extract analysis metadata from a chunk."""
    # Get chunk text
    text = chunk.text if hasattr(chunk, 'text') else str(chunk)

    # Detect if chunk contains tables
    has_table = False
    table_count = 0

    if hasattr(chunk, 'meta') and chunk.meta:
        # Check if this chunk is a table or contains table references
        if hasattr(chunk.meta, 'doc_items'):
            for item in chunk.meta.doc_items:
                if hasattr(item, 'label') and item.label == DocItemLabel.TABLE:
                    has_table = True
                    table_count += 1

    # Check text content for table markers
    if '|' in text or 'Table' in text or 'table' in text:
        has_table = True

    return {
        'chunk_num': chunk_num,
        'char_count': len(text),
        'line_count': text.count('\n') + 1,
        'has_table': has_table,
        'table_count': table_count,
        'preview': text[:200] + '...' if len(text) > 200 else text,
        'full_text': text,
    }


def generate_markdown_report(strategy_name: str, chunks: list[BaseChunk], output_path: Path):
    """Generate a markdown report showing chunk boundaries and analysis."""

    analyses = [analyze_chunk(chunk, i+1) for i, chunk in enumerate(chunks)]

    # Write report
    with output_path.open('w') as f:
        f.write(f"# {strategy_name} - Chunking Evaluation\n\n")
        f.write(f"**Document**: TMUS_992_Q423.pdf\n\n")
        f.write(f"**Total Chunks**: {len(chunks)}\n\n")

        # Summary statistics
        total_chars = sum(a['char_count'] for a in analyses)
        chunks_with_tables = sum(1 for a in analyses if a['has_table'])
        total_tables = sum(a['table_count'] for a in analyses)

        f.write("## Summary Statistics\n\n")
        f.write(f"- **Total Characters**: {total_chars:,}\n")
        f.write(f"- **Chunks with Tables**: {chunks_with_tables}\n")
        f.write(f"- **Total Tables Detected**: {total_tables}\n")
        f.write(f"- **Average Chunk Size**: {total_chars // len(chunks):,} chars\n\n")

        # Chunk details
        f.write("## Chunk Breakdown\n\n")

        for analysis in analyses:
            f.write(f"### Chunk {analysis['chunk_num']}\n\n")
            f.write(f"- **Size**: {analysis['char_count']} chars, {analysis['line_count']} lines\n")
            f.write(f"- **Contains Table**: {'✓ Yes' if analysis['has_table'] else '✗ No'}\n")
            if analysis['table_count'] > 0:
                f.write(f"- **Table Count**: {analysis['table_count']}\n")
            f.write(f"\n**Preview**:\n```\n{analysis['preview']}\n```\n\n")

            f.write("<details>\n<summary>Full Chunk Content</summary>\n\n")
            f.write(f"```\n{analysis['full_text']}\n```\n\n")
            f.write("</details>\n\n")
            f.write("---\n\n")

    logger.info(f"Report written to {output_path}")


def main():
    """Main evaluation workflow."""

    # Input document
    doc_path = Path("documents/TMUS_992_Q423.pdf")
    if not doc_path.exists():
        logger.error(f"Document not found: {doc_path}")
        sys.exit(1)

    # Output directory
    output_dir = Path("output/chunking_evaluation")
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info(f"Converting document: {doc_path}")
    converter = setup_converter()
    result = converter.convert(doc_path)
    doc = result.document

    logger.info(f"Document converted successfully. Pages: {len(doc.pages)}")

    # Evaluate each chunking strategy
    strategies = [
        ("Base Chunker", chunk_with_base),
        ("Hierarchical Chunker", chunk_with_hierarchical),
        ("Hybrid Chunker", chunk_with_hybrid),
    ]

    for strategy_name, chunk_func in strategies:
        logger.info(f"\n{'=' * 60}")
        logger.info(f"Evaluating: {strategy_name}")
        logger.info(f"{'=' * 60}")

        try:
            chunks = chunk_func(doc)

            if chunks:
                report_path = output_dir / f"{strategy_name.lower().replace(' ', '_')}_report.md"
                generate_markdown_report(strategy_name, chunks, report_path)
            else:
                logger.warning(f"No chunks generated for {strategy_name}")

        except Exception as e:
            logger.error(f"Error in {strategy_name}: {e}", exc_info=True)

    logger.info(f"\n{'=' * 60}")
    logger.info("Evaluation complete!")
    logger.info(f"Reports available in: {output_dir}")
    logger.info(f"{'=' * 60}")


if __name__ == "__main__":
    main()