#!/usr/bin/env python3
"""Test LangChain MarkdownHeaderTextSplitter for chunking by sheet."""

import sys
import subprocess
import tempfile
from pathlib import Path
from collections import defaultdict

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

from docling.document_converter import DocumentConverter
from openpyxl import load_workbook
from langchain_text_splitters import MarkdownHeaderTextSplitter


def convert_xls(xls_path: str) -> str:
    """Convert .xls to .xlsx."""
    temp_dir = tempfile.mkdtemp()
    result = subprocess.run([
        '/Applications/LibreOffice.app/Contents/MacOS/soffice',
        '--headless',
        '--convert-to', 'xlsx',
        '--outdir', temp_dir,
        xls_path
    ], capture_output=True, text=True, timeout=60)

    xlsx_files = list(Path(temp_dir).glob('*.xlsx'))
    return str(xlsx_files[0]) if xlsx_files else None


def get_sheet_names(xlsx_path: str):
    """Get sheet names from Excel file."""
    wb = load_workbook(xlsx_path, read_only=True)
    names = wb.sheetnames
    wb.close()
    return names


def chunk_by_sheet_with_items(doc, sheet_names: list):
    """Chunk by sheet using doc.iterate_items() grouped by page_no."""

    # Group items by page_no
    sheets = defaultdict(list)

    for item_tuple in doc.iterate_items():
        item = item_tuple[0] if isinstance(item_tuple, tuple) else item_tuple

        page_no = None
        if hasattr(item, 'prov') and item.prov:
            for prov in item.prov:
                if hasattr(prov, 'page_no'):
                    page_no = prov.page_no
                    break

        if page_no is not None:
            sheets[page_no].append(item)

    # Create chunks per sheet
    chunks = []

    for page_no in sorted(sheets.keys()):
        sheet_idx = page_no - 1
        sheet_name = sheet_names[sheet_idx] if 0 <= sheet_idx < len(sheet_names) else f"Sheet{page_no}"

        items = sheets[page_no]

        # Export items to markdown
        texts = []
        for item in items:
            if hasattr(item, 'export_to_markdown'):
                text = item.export_to_markdown()
                if text and text.strip():
                    texts.append(text.strip())

        combined_text = "\n\n".join(texts)

        if combined_text.strip():
            chunks.append({
                'sheet_name': sheet_name,
                'page_no': page_no,
                'text': combined_text,
                'metadata': {
                    'sheet_name': sheet_name,
                    'sheet_number': page_no,
                }
            })

    return chunks


def chunk_with_langchain(markdown: str, max_chunk_size: int = 4000):
    """Chunk markdown using LangChain MarkdownHeaderTextSplitter."""

    # Define headers to split on (H1, H2, H3)
    headers_to_split_on = [
        ("#", "Header 1"),
        ("##", "Header 2"),
        ("###", "Header 3"),
    ]

    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=headers_to_split_on,
        strip_headers=False
    )

    chunks = markdown_splitter.split_text(markdown)

    return chunks


def main(file_path: str):
    """Test both chunking approaches."""

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

    # Convert if needed
    if file_path.endswith('.xls'):
        xlsx_path = convert_xls(file_path)
        print(f"Converted to: {xlsx_path}\n")
    else:
        xlsx_path = file_path

    # Get sheet names
    sheet_names = get_sheet_names(xlsx_path)
    print(f"Sheet names ({len(sheet_names)}): {sheet_names}\n")

    # Convert with Docling
    print("Converting with Docling...")
    converter = DocumentConverter()
    result = converter.convert(Path(xlsx_path))
    doc = result.document

    # Export to markdown
    markdown = doc.export_to_markdown()
    print(f"Exported to markdown: {len(markdown)} chars\n")

    # APPROACH 1: Chunk by sheet using iterate_items()
    print(f"{'='*80}")
    print("APPROACH 1: Chunk by Sheet (using iterate_items)")
    print(f"{'='*80}\n")

    sheet_chunks = chunk_by_sheet_with_items(doc, sheet_names)
    print(f"Created {len(sheet_chunks)} chunks (one per sheet with content)\n")

    for chunk in sheet_chunks:
        print(f"  {chunk['sheet_name']}: {len(chunk['text'])} chars")

    # Check target sheets
    print(f"\nTarget sheets:")
    targets = ['ステータス戻し実施可否判定', 'チェック対象判定']
    for target in targets:
        found = [c for c in sheet_chunks if c['sheet_name'] == target]
        if found:
            print(f"  ✓ {target}: {len(found[0]['text'])} chars")
        else:
            print(f"  ✗ {target}: NOT FOUND")

    # APPROACH 2: LangChain MarkdownHeaderTextSplitter
    print(f"\n{'='*80}")
    print("APPROACH 2: LangChain MarkdownHeaderTextSplitter")
    print(f"{'='*80}\n")

    langchain_chunks = chunk_with_langchain(markdown)
    print(f"Created {len(langchain_chunks)} chunks\n")

    for idx, chunk in enumerate(langchain_chunks[:10], 1):
        content_preview = chunk.page_content[:100] if len(chunk.page_content) > 100 else chunk.page_content
        metadata = chunk.metadata if hasattr(chunk, 'metadata') else {}
        print(f"  Chunk {idx}: {len(chunk.page_content)} chars")
        print(f"    Metadata: {metadata}")
        print(f"    Preview: {content_preview}...")
        print()

    # RECOMMENDED: Combine both approaches
    # Use iterate_items to chunk by sheet, then use RecursiveCharacterTextSplitter if needed
    print(f"{'='*80}")
    print("RECOMMENDED APPROACH:")
    print(f"{'='*80}\n")
    print("1. Use doc.iterate_items() to group content by sheet (page_no)")
    print("2. For each sheet, export items to markdown")
    print("3. If sheet content > max_chunk_size, use RecursiveCharacterTextSplitter")
    print("4. This ensures:")
    print("   - All sheets are included (no content lost)")
    print("   - Sheet metadata preserved")
    print("   - Large sheets are split intelligently")


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

    main(sys.argv[1])