#!/usr/bin/env python3
"""Test refactored chunking using export_to_markdown(page_no=X) and LangChain."""

import sys
import subprocess
import tempfile
from pathlib import Path

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 RecursiveCharacterTextSplitter


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_refactored(doc, sheet_names: list, max_characters: int = 4000):
    """Chunk by sheet using export_to_markdown(page_no=X) and LangChain splitter."""

    chunks = []

    # Process each sheet (page_no is 1-based)
    for page_no in range(1, len(sheet_names) + 1):
        sheet_idx = page_no - 1  # Convert 1-based to 0-based
        sheet_name = sheet_names[sheet_idx]

        # Export this sheet's markdown using page_no parameter
        try:
            sheet_markdown = doc.export_to_markdown(page_no=page_no)
        except Exception as e:
            print(f"Failed to export sheet {page_no} ({sheet_name}): {e}")
            continue

        if not sheet_markdown or not sheet_markdown.strip():
            print(f"Sheet {page_no} ({sheet_name}) has no content, skipping")
            continue

        # Check if we need to split this sheet
        if len(sheet_markdown) <= max_characters:
            # Sheet fits in one chunk
            chunks.append({
                'sheet_name': sheet_name,
                'page_no': page_no,
                'text': sheet_markdown.strip(),
                'chunk_index': 0,
                'total_chunks': 1,
            })

            print(f"✓ Created chunk for sheet '{sheet_name}' (page_no={page_no}): {len(sheet_markdown)} chars")
        else:
            # Sheet is too large, split it using RecursiveCharacterTextSplitter
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=max_characters,
                chunk_overlap=200,
                length_function=len,
                separators=["\n\n", "\n", " ", ""]
            )

            split_texts = text_splitter.split_text(sheet_markdown)

            print(
                f"✓ Sheet '{sheet_name}' (page_no={page_no}) is large ({len(sheet_markdown)} chars), "
                f"split into {len(split_texts)} chunks"
            )

            for idx, split_text in enumerate(split_texts):
                chunks.append({
                    'sheet_name': sheet_name,
                    'page_no': page_no,
                    'text': split_text.strip(),
                    'chunk_index': idx,
                    'total_chunks': len(split_texts),
                })

    return chunks


def main(file_path: str):
    """Test refactored chunking approach."""

    print(f"Testing refactored chunking 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

    print(f"Converted successfully\n")

    # Test refactored chunking
    print(f"{'='*80}")
    print("REFACTORED CHUNKING (export_to_markdown(page_no=X) + LangChain)")
    print(f"{'='*80}\n")

    chunks = chunk_by_sheet_refactored(doc, sheet_names)

    print(f"\nCreated {len(chunks)} total chunks from {len(sheet_names)} sheets\n")

    # Show chunk summary
    print(f"{'='*80}")
    print("CHUNK SUMMARY:")
    print(f"{'='*80}\n")

    for chunk in chunks:
        chunk_label = f"{chunk['chunk_index']+1}/{chunk['total_chunks']}" if chunk['total_chunks'] > 1 else "1/1"
        print(f"  {chunk['sheet_name']} [{chunk_label}]: {len(chunk['text'])} chars")

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

    target_sheets = ['ステータス戻し実施可否判定', 'チェック対象判定']
    for target in target_sheets:
        found = [c for c in chunks if c['sheet_name'] == target]
        if found:
            total_chars = sum(len(c['text']) for c in found)
            print(f"✓ {target}:")
            print(f"    Chunks: {len(found)}")
            print(f"    Total chars: {total_chars}")
            print(f"    Preview: {found[0]['text'][:200]}...")
        else:
            print(f"✗ {target}: NOT FOUND")


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

    main(sys.argv[1])