#!/usr/bin/env python3
"""Test chunking by sheet using Docling document items."""

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


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(doc, sheet_names: list):
    """Chunk document by sheet using iterate_items()."""

    # Group items by page_no (sheet number)
    sheets = defaultdict(list)

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

        # Get page_no from provenance
        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)

    print(f"Found items in {len(sheets)} sheets (by page_no)")
    print(f"Sheet page_no values: {sorted(sheets.keys())}\n")

    # Create chunks per sheet
    chunks_by_sheet = {}

    for page_no in sorted(sheets.keys()):
        sheet_idx = page_no - 1  # Convert 1-based to 0-based
        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:
            # Try to get text from item
            if hasattr(item, 'export_to_markdown'):
                text = item.export_to_markdown()
                if text and text.strip():
                    texts.append(text.strip())
            elif hasattr(item, '__str__'):
                text = str(item)
                if text and text.strip() and not text.startswith('<'):
                    texts.append(text.strip())

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

        chunks_by_sheet[sheet_name] = {
            'page_no': page_no,
            'item_count': len(items),
            'text': combined_text,
            'text_length': len(combined_text)
        }

    return chunks_by_sheet


def main(file_path: str):
    """Test chunking by sheet."""

    print(f"Testing chunk-by-sheet 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
    doc.export_to_markdown

    print(f"Converted successfully\n")

    # Chunk by sheet
    print(f"{'='*80}")
    print("CHUNKING BY SHEET:")
    print(f"{'='*80}\n")

    chunks = chunk_by_sheet(doc, sheet_names)

    print(f"Created {len(chunks)} chunks (one per sheet)\n")

    # Show results
    for sheet_name, chunk_info in chunks.items():
        print(f"\n{'-'*80}")
        print(f"Sheet: {sheet_name}")
        print(f"  Page number: {chunk_info['page_no']}")
        print(f"  Items: {chunk_info['item_count']}")
        print(f"  Text length: {chunk_info['text_length']} chars")
        print(f"  Text preview: {chunk_info['text'][:200]}...")

    # Check target sheets
    print(f"\n{'='*80}")
    print("TARGET SHEETS CHECK:")
    print(f"{'='*80}\n")

    target_sheets = ['ステータス戻し実施可否判定', 'チェック対象判定']
    for target in target_sheets:
        if target in chunks:
            info = chunks[target]
            print(f"✓ {target}:")
            print(f"    Text length: {info['text_length']} chars")
            print(f"    Items: {info['item_count']}")
            if info['text_length'] > 0:
                print(f"    HAS TEXT!")
            else:
                print(f"    EMPTY!")
        else:
            print(f"✗ {target}: NOT FOUND")


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

    main(sys.argv[1])