#!/usr/bin/env python3
"""
Test script for Docling + LangChain semantic chunking pipeline.

This script demonstrates:
1. Reading document files (.docx, .doc, .docm) with Docling
2. Converting to Markdown
3. Splitting with LangChain RecursiveCharacterTextSplitter
4. Printing results with detailed analysis
"""

import sys
from pathlib import Path
from typing import List

# Import Docling for document conversion
from docling.document_converter import DocumentConverter

# Import LangChain for semantic chunking
from langchain.text_splitter import RecursiveCharacterTextSplitter


def read_document_with_docling(file_path: str) -> str:
    """
    Read a document file and convert to Markdown using Docling.

    Args:
        file_path: Path to the document (.docx, .doc, .docm)

    Returns:
        Markdown string of the document content
    """
    print(f"\n{'='*80}")
    print(f"📄 Reading document: {Path(file_path).name}")
    print(f"{'='*80}\n")

    # Initialize Docling converter
    converter = DocumentConverter()

    # Convert document to Docling format
    result = converter.convert(file_path)

    if result and result.document:
        # Export to Markdown
        markdown = result.document.export_to_markdown()

        print(f"✅ Successfully converted to Markdown")
        print(f"📊 Document stats:")
        print(f"   - Characters: {len(markdown):,}")
        print(f"   - Words: {len(markdown.split()):,}")
        print(f"   - Lines: {len(markdown.splitlines()):,}")
        print(f"\n{'─'*80}\n")

        return markdown
    else:
        raise RuntimeError(f"Failed to convert document: {file_path}")


def split_with_langchain(
    markdown: str,
    chunk_size: int = 1000,
    chunk_overlap: int = 200
) -> List[str]:
    """
    Split markdown content using LangChain's semantic splitter.

    Args:
        markdown: Markdown content from Docling
        chunk_size: Maximum characters per chunk (default: 1000)
        chunk_overlap: Character overlap between chunks (default: 200)

    Returns:
        List of text chunks
    """
    print(f"🔪 Splitting document with LangChain")
    print(f"⚙️  Configuration:")
    print(f"   - Chunk size: {chunk_size} characters")
    print(f"   - Overlap: {chunk_overlap} characters")
    print(f"   - Strategy: RecursiveCharacterTextSplitter\n")

    # Initialize LangChain splitter with semantic separators
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=[
            "\n\n",      # Paragraph boundaries (highest priority)
            "\n",        # Line breaks
            ". ",        # Sentence endings
            "? ",        # Question endings
            "! ",        # Exclamation endings
            "; ",        # Clause boundaries
            ", ",        # Comma boundaries
            " ",         # Word boundaries
            ""           # Character-level (fallback)
        ],
        is_separator_regex=False
    )

    # Split the text
    chunks = text_splitter.split_text(markdown)

    print(f"✅ Splitting complete")
    print(f"📊 Chunking stats:")
    print(f"   - Total chunks: {len(chunks)}")

    if chunks:
        avg_size = sum(len(c) for c in chunks) // len(chunks)
        min_size = min(len(c) for c in chunks)
        max_size = max(len(c) for c in chunks)
        print(f"   - Average chunk size: {avg_size} characters")
        print(f"   - Min chunk size: {min_size} characters")
        print(f"   - Max chunk size: {max_size} characters")

    print(f"\n{'─'*80}\n")

    return chunks


def print_chunks(chunks: List[str], show_full_content: bool = False):
    """
    Print the chunks with detailed information.

    Args:
        chunks: List of text chunks
        show_full_content: If True, show full chunk content; if False, show preview
    """
    print(f"📑 Chunk Details:")
    print(f"{'='*80}\n")

    for idx, chunk in enumerate(chunks):
        print(f"Chunk #{idx + 1} of {len(chunks)}")
        print(f"{'─'*80}")
        print(f"📏 Size: {len(chunk)} characters, {len(chunk.split())} words")
        print(f"📄 Lines: {len(chunk.splitlines())}")

        # Check if chunk contains markdown tables
        has_table = '|' in chunk and '---' in chunk
        if has_table:
            print(f"📊 Contains markdown table")

        # Check for headings
        headings = [line for line in chunk.splitlines() if line.startswith('#')]
        if headings:
            print(f"📌 Headings: {len(headings)}")

        print(f"\n📝 Content:")
        print(f"{'─'*80}")

        if show_full_content:
            # Show full content
            print(chunk)
        else:
            # Show preview (first 300 chars)
            preview_length = min(300, len(chunk))
            preview = chunk[:preview_length]
            if len(chunk) > preview_length:
                preview += "\n\n[... truncated ...]"
            print(preview)

        print(f"\n{'='*80}\n")


def analyze_overlap(chunks: List[str], chunk_overlap: int):
    """
    Analyze the actual overlap between consecutive chunks.

    Args:
        chunks: List of text chunks
        chunk_overlap: Expected overlap size
    """
    if len(chunks) < 2:
        print("⚠️  Not enough chunks to analyze overlap")
        return

    print(f"🔍 Overlap Analysis:")
    print(f"{'='*80}\n")

    for i in range(len(chunks) - 1):
        current_chunk = chunks[i]
        next_chunk = chunks[i + 1]

        # Find the overlap by checking the end of current chunk with start of next
        overlap_found = 0
        max_check = min(len(current_chunk), len(next_chunk), chunk_overlap + 100)

        for overlap_size in range(max_check, 0, -1):
            if current_chunk[-overlap_size:] == next_chunk[:overlap_size]:
                overlap_found = overlap_size
                break

        print(f"Chunk {i + 1} → Chunk {i + 2}:")
        print(f"   - Overlap found: {overlap_found} characters")

        if overlap_found > 0:
            overlap_text = current_chunk[-overlap_found:][:50]  # First 50 chars of overlap
            if len(overlap_text) < overlap_found:
                overlap_text += "..."
            print(f"   - Overlap text: \"{overlap_text}\"")

        print()


def main():
    """Main execution function."""
    # Check command line arguments
    if len(sys.argv) < 2:
        print("Usage: python test_langchain_splitter.py <document_path> [--full]")
        print("\nExample:")
        print("  python test_langchain_splitter.py CustomerDocument/sample.docx")
        print("  python test_langchain_splitter.py CustomerDocument/sample.docx --full")
        print("\nOptions:")
        print("  --full    Show full content of each chunk (default: show preview only)")
        sys.exit(1)

    file_path = sys.argv[1]
    show_full_content = "--full" in sys.argv

    # Validate file exists
    if not Path(file_path).exists():
        print(f"❌ Error: File not found: {file_path}")
        sys.exit(1)

    try:
        # Step 1: Read document with Docling
        markdown = read_document_with_docling(file_path)

        # Optional: Print markdown preview
        print("📄 Markdown Preview (first 500 characters):")
        print(f"{'─'*80}")
        print(markdown[:500])
        if len(markdown) > 500:
            print("\n[... truncated ...]")
        print(f"\n{'='*80}\n")

        # Step 2: Split with LangChain
        chunks = split_with_langchain(
            markdown,
            chunk_size=1000,
            chunk_overlap=200
        )

        # Step 3: Print chunks
        print_chunks(chunks, show_full_content=show_full_content)

        # Step 4: Analyze overlap
        analyze_overlap(chunks, chunk_overlap=200)

        # Summary
        print(f"\n{'='*80}")
        print(f"✅ Pipeline Complete!")
        print(f"{'='*80}")
        print(f"📊 Summary:")
        print(f"   - Document: {Path(file_path).name}")
        print(f"   - Original size: {len(markdown):,} characters")
        print(f"   - Total chunks: {len(chunks)}")
        print(f"   - Avg chunk size: {sum(len(c) for c in chunks) // len(chunks) if chunks else 0} characters")
        print(f"{'='*80}\n")

    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()