#!/usr/bin/env python3
"""Test script for large table parsing strategy.

This script validates an alternative approach for parsing very large Excel tables
that exceed Docling's processing limits:

Strategy:
1. Use border detection to identify table boundaries and sizes
2. If table has > MAX_ROWS rows, use chunked processing:
   a. Extract first 10 rows to separate Excel file
   b. Convert header-only file to HTML with Docling (preserves multi-level headers)
   c. Parse HTML to extract and normalize headers using header_extractor
   d. Iterate through data rows in chunks
   e. Map each data row to normalized headers
   f. Export to markdown/JSON

Usage:
    python tests/manual/test_large_table_strategy.py [file_path]

Example:
    python tests/manual/test_large_table_strategy.py "CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.2 JP1шинхоЪщаЕчЫо.xlsx"
"""

import sys
import time
import tempfile
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from openpyxl import Workbook
from openpyxl.utils import get_column_letter, column_index_from_string

from src.extraction_v2.detect_border import BorderTableDetector
from src.extraction_v2.excel_to_html_converter import ExcelToHtmlConverter
from src.extraction_v2.excel_loader import load_excel_file
from src.extraction.header_extractor import HeaderExtractor
from src.extraction.models import TableBoundary


# Configuration
MAX_ROWS_FOR_DOCLING = 1000  # Tables larger than this use chunked processing
HEADER_SAMPLE_ROWS = 10      # Number of rows to extract for header detection
CHUNK_SIZE = 100             # Number of data rows to process at once


def print_header(title: str):
    """Print a formatted section header."""
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80)


def print_step(step_num: int, total_steps: int, description: str):
    """Print a step indicator."""
    print(f"\n[Step {step_num}/{total_steps}] {description}")
    print("-" * 80)


def detect_tables(file_path: str) -> List[Dict[str, Any]]:
    """
    Detect all tables in the Excel file using border detection.

    Args:
        file_path: Path to Excel file

    Returns:
        List of table information dictionaries
    """
    print_step(1, 7, "Detecting Tables Using Border Detection")

    detector = BorderTableDetector()
    tables = detector.detect_tables_by_border(file_path)

    print(f"\n✅ Found {len(tables)} table(s)")

    for idx, table in enumerate(tables, 1):
        print(f"\nTable {idx}:")
        print(f"  Sheet:      {table['sheet']}")
        print(f"  Range:      {table['range']}")
        print(f"  Rows:       {table['start_row']} → {table['end_row']}")
        print(f"  Columns:    {table['start_col']} → {table['end_col']}")
        print(f"  Cell count: {table['cell_count']:,}")

        num_rows = table['end_row'] - table['start_row'] + 1
        if num_rows > MAX_ROWS_FOR_DOCLING:
            print(f"  ⚠️  LARGE TABLE: {num_rows:,} rows (exceeds {MAX_ROWS_FOR_DOCLING} row limit)")
        else:
            print(f"  ✓  Normal size: {num_rows:,} rows")

    return tables


def extract_header_rows(file_path: str, table: Dict[str, Any]) -> str:
    """
    Extract first N rows of a table to a temporary Excel file.

    Args:
        file_path: Path to source Excel file
        table: Table information dictionary

    Returns:
        Path to temporary Excel file with header rows
    """
    print_step(2, 7, f"Extracting First {HEADER_SAMPLE_ROWS} Rows for Header Detection")

    # Load source workbook
    source_wb = load_excel_file(file_path, data_only=False)
    source_sheet = source_wb[table['sheet']]

    # Create new workbook with just header rows
    dest_wb = Workbook()
    dest_sheet = dest_wb.active
    dest_sheet.title = table['sheet']

    # Get column range
    start_col_idx = column_index_from_string(table['start_col'])
    end_col_idx = column_index_from_string(table['end_col'])

    # Calculate how many rows to extract
    num_rows = min(HEADER_SAMPLE_ROWS, table['end_row'] - table['start_row'] + 1)

    print(f"   Extracting rows {table['start_row']} to {table['start_row'] + num_rows - 1}")
    print(f"   Columns: {table['start_col']} to {table['end_col']}")

    # Copy header rows
    for row_offset in range(num_rows):
        source_row = table['start_row'] + row_offset
        dest_row = row_offset + 1

        for col_offset in range(end_col_idx - start_col_idx + 1):
            source_col = start_col_idx + col_offset
            dest_col = col_offset + 1

            source_cell = source_sheet.cell(row=source_row, column=source_col)
            dest_cell = dest_sheet.cell(row=dest_row, column=dest_col)

            # Copy value and basic style
            dest_cell.value = source_cell.value
            if source_cell.has_style:
                dest_cell.font = source_cell.font.copy()
                dest_cell.border = source_cell.border.copy()
                dest_cell.fill = source_cell.fill.copy()
                dest_cell.alignment = source_cell.alignment.copy()

    # Copy merged cells within the header range
    for merged_range in source_sheet.merged_cells.ranges:
        # Check if merged range is within our extracted area
        if (merged_range.min_row >= table['start_row'] and
            merged_range.max_row < table['start_row'] + num_rows and
            merged_range.min_col >= start_col_idx and
            merged_range.max_col <= end_col_idx):

            # Adjust to dest coordinates
            new_min_row = merged_range.min_row - table['start_row'] + 1
            new_max_row = merged_range.max_row - table['start_row'] + 1
            new_min_col = merged_range.min_col - start_col_idx + 1
            new_max_col = merged_range.max_col - start_col_idx + 1

            dest_sheet.merge_cells(
                start_row=new_min_row,
                start_column=new_min_col,
                end_row=new_max_row,
                end_column=new_max_col
            )

    # Save to temporary file
    temp_file = tempfile.NamedTemporaryFile(suffix='.xlsx', delete=False)
    temp_path = temp_file.name
    temp_file.close()

    dest_wb.save(temp_path)
    dest_wb.close()
    source_wb.close()

    print(f"   ✅ Saved header rows to: {temp_path}")
    print(f"   File size: {Path(temp_path).stat().st_size:,} bytes")

    return temp_path


def convert_headers_to_html(header_file: str) -> str:
    """
    Convert header-only Excel file to HTML using Docling.

    Args:
        header_file: Path to temporary Excel file with headers

    Returns:
        HTML content string
    """
    print_step(3, 7, "Converting Header Rows to HTML Using Docling")

    converter = ExcelToHtmlConverter()
    html = converter.convert_to_html(header_file)

    if html:
        print(f"   ✅ HTML generated: {len(html):,} characters")
        print(f"   Tables in HTML: {html.count('<table')}")
    else:
        print(f"   ❌ HTML conversion failed")

    return html


def extract_normalized_headers(header_file: str, table: Dict[str, Any]) -> Dict[str, Any]:
    """
    Extract and normalize multi-level headers from header-only Excel file.

    Args:
        header_file: Path to temporary Excel file with header rows
        table: Table information dictionary

    Returns:
        Dictionary with:
        - 'column_mapping': dict mapping column_letter -> normalized_header_name
        - 'boundary': TableBoundary object
        - 'header_info': dict of HeaderInfo objects
    """
    print_step(4, 7, "Extracting and Normalizing Multi-Level Headers")

    # Load the header-only file
    workbook = load_excel_file(header_file, data_only=False)
    sheet = workbook[table['sheet']]

    # Create table boundary for the header file
    boundary = TableBoundary(
        sheet=table['sheet'],
        table_id=1,
        start_row=1,  # Header file starts at row 1
        end_row=HEADER_SAMPLE_ROWS,
        start_col=table['start_col'],
        end_col=table['end_col']
    )

    print(f"   Table boundary: {boundary.start_col}{boundary.start_row}:{boundary.end_col}{boundary.end_row}")
    print(f"   Extracting headers using HeaderExtractor...")

    # Use HeaderExtractor to extract normalized headers
    extractor = HeaderExtractor()
    try:
        headers_info = extractor.extract_headers(sheet, boundary)
        
        # Detect actual header depth
        header_depth = extractor._detect_header_depth(sheet, boundary)

        print(f"   ✅ Extracted {len(headers_info)} column headers")
        print(f"   ✅ Detected header depth: {header_depth} rows")

        # Create column mapping: column_letter -> header display name
        # Remove table name prefix (first level) if present
        column_mapping = {}
        for col_letter, header_info in headers_info.items():
            # Get the original header with all levels
            full_header = header_info.display

            # Split by " > " to get individual levels
            levels = full_header.split(" > ")
            
            # Remove consecutive duplicates (caused by vertically merged cells)
            deduplicated_levels = []
            for level in levels:
                if not deduplicated_levels or level != deduplicated_levels[-1]:
                    deduplicated_levels.append(level)
            levels = deduplicated_levels
            
            # Remove first level (table/sheet name) if there are multiple levels
            if len(levels) > 1:
                levels = levels[1:]
            
            # Rejoin levels
            header_name = " > ".join(levels) if levels else full_header
            column_mapping[col_letter] = header_name

        # Show sample headers
        sample_cols = list(column_mapping.items())[:5]
        print(f"   Sample headers:")
        for col_letter, header_name in sample_cols:
            print(f"     {col_letter}: {header_name[:60]}{'...' if len(header_name) > 60 else ''}")

        workbook.close()

        return {
            'column_mapping': column_mapping,
            'boundary': boundary,
            'header_info': headers_info,
            'header_depth': header_depth
        }

    except Exception as e:
        print(f"   ⚠️  Header extraction failed: {e}")
        print(f"   Using column letters as fallback")

        # Fallback: use column letters as headers
        start_col_idx = column_index_from_string(table['start_col'])
        end_col_idx = column_index_from_string(table['end_col'])
        column_mapping = {}
        for col_idx in range(start_col_idx, end_col_idx + 1):
            col_letter = get_column_letter(col_idx)
            column_mapping[col_letter] = col_letter

        workbook.close()

        return {
            'column_mapping': column_mapping,
            'boundary': boundary,
            'header_info': {},
            'header_depth': HEADER_SAMPLE_ROWS  # Fallback to full sample
        }


def process_data_rows(file_path: str, table: Dict[str, Any], headers: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Process data rows in chunks and map to normalized headers.

    Args:
        file_path: Path to source Excel file
        table: Table information dictionary
        headers: Normalized header information

    Returns:
        List of row dictionaries
    """
    print_step(5, 7, "Processing Data Rows in Chunks")

    workbook = load_excel_file(file_path, read_only=True, data_only=True)
    sheet = workbook[table['sheet']]

    # Calculate data row range using actual header depth
    actual_header_depth = headers.get('header_depth', HEADER_SAMPLE_ROWS)
    data_start_row = table['start_row'] + actual_header_depth
    data_end_row = table['end_row']
    total_data_rows = data_end_row - data_start_row + 1

    print(f"   Header depth: {actual_header_depth} rows")
    print(f"   Data rows: {data_start_row} to {data_end_row} ({total_data_rows:,} rows)")

    # Get column range
    start_col_idx = column_index_from_string(table['start_col'])
    end_col_idx = column_index_from_string(table['end_col'])
    num_cols = end_col_idx - start_col_idx + 1

    print(f"   Columns: {table['start_col']} to {table['end_col']} ({num_cols} columns)")
    print(f"   Processing in chunks of {CHUNK_SIZE} rows...")

    rows_data = []
    rows_processed = 0

    # Process in chunks using iter_rows (much faster than cell-by-cell access)
    for chunk_start in range(data_start_row, data_end_row + 1, CHUNK_SIZE):
        chunk_end = min(chunk_start + CHUNK_SIZE - 1, data_end_row)
        chunk_size = chunk_end - chunk_start + 1

        print(f"   Processing rows {chunk_start} to {chunk_end} ({chunk_size} rows)...")

        # Read chunk using iter_rows (FAST!)
        for row in sheet.iter_rows(
            min_row=chunk_start,
            max_row=chunk_end,
            min_col=start_col_idx,
            max_col=end_col_idx,
            values_only=True  # Only get values, not Cell objects
        ):
            # Skip empty rows (all values are None or empty)
            if all(value is None or (isinstance(value, str) and not value.strip()) for value in row):
                continue
            
            row_data = {}

            # Map each cell value to its normalized header name
            column_mapping = headers.get('column_mapping', {})
            for col_offset, value in enumerate(row):
                col_idx = start_col_idx + col_offset
                col_letter = get_column_letter(col_idx)

                # Use normalized header name if available, fallback to column letter
                header_name = column_mapping.get(col_letter, col_letter)
                row_data[header_name] = value

            rows_data.append(row_data)
            rows_processed += 1

        # Progress update
        progress = (rows_processed / total_data_rows) * 100
        print(f"      Progress: {rows_processed:,}/{total_data_rows:,} rows ({progress:.1f}%)")

        # Early exit for testing (process only first 2 chunks)
        if len(rows_data) >= CHUNK_SIZE * 2:
            print(f"   ⚠️  Early exit for testing (processed {len(rows_data)} rows)")
            break

    workbook.close()

    print(f"   ✅ Processed {len(rows_data):,} data rows")

    return rows_data


def export_results(table: Dict[str, Any], headers: Dict[str, Any], data: List[Dict[str, Any]]) -> str:
    """
    Export parsed results to markdown and JSON formats.

    Args:
        table: Table information
        headers: Header information
        data: List of data rows

    Returns:
        Path to output markdown file
    """
    print_step(6, 7, "Exporting Results to Markdown and JSON")

    # Create output directory
    output_dir = Path("./tmp")
    output_dir.mkdir(exist_ok=True)

    # Generate output filename
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    safe_sheet_name = table['sheet'].replace('/', '_').replace('\\', '_')
    base_filename = f"large_table_{safe_sheet_name}_{timestamp}"

    md_path = output_dir / f"{base_filename}.md"
    json_path = output_dir / f"{base_filename}.json"

    # Build markdown content
    lines = []
    lines.append(f"# Large Table Extraction Results\n")
    lines.append(f"**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
    lines.append(f"## Table Information\n")
    lines.append(f"- **Sheet:** {table['sheet']}")
    lines.append(f"- **Range:** {table['range']}")
    lines.append(f"- **Total Rows:** {table['end_row'] - table['start_row'] + 1:,}")
    lines.append(f"- **Columns:** {table['start_col']} to {table['end_col']}")
    lines.append(f"- **Data Rows Processed:** {len(data):,}\n")

    lines.append(f"## Sample Data (First 10 Rows)\n")

    # Show first 10 rows
    import json
    for idx, row in enumerate(data[:10], 1):
        lines.append(f"### Row {idx}\n")
        lines.append("```json")
        lines.append(json.dumps(row, indent=2, ensure_ascii=False))
        lines.append("```\n")

    # Save markdown file
    with open(md_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(lines))

    print(f"   ✅ Markdown saved to: {md_path}")
    print(f"      File size: {md_path.stat().st_size:,} bytes")

    # Export JSON file with all data
    json_output = {
        "metadata": {
            "generated": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            "sheet": table['sheet'],
            "range": table['range'],
            "total_rows": table['end_row'] - table['start_row'] + 1,
            "columns": {
                "start": table['start_col'],
                "end": table['end_col'],
                "count": len(headers.get('column_mapping', {}))
            },
            "data_rows_processed": len(data)
        },
        "headers": headers.get('column_mapping', {}),
        "data": data
    }

    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(json_output, f, indent=2, ensure_ascii=False)

    print(f"   ✅ JSON saved to: {json_path}")
    print(f"      File size: {json_path.stat().st_size:,} bytes")
    print(f"      Contains {len(data)} data rows")

    return str(md_path)


def test_large_table_strategy(file_path: str):
    """
    Test the large table parsing strategy.

    Args:
        file_path: Path to Excel file to test
    """
    print_header("LARGE TABLE PARSING STRATEGY TEST")

    test_file = Path(file_path)

    if not test_file.exists():
        print(f"❌ Error: File not found: {test_file}")
        return

    print(f"\n📁 Test File: {test_file}")
    print(f"   Size: {test_file.stat().st_size / (1024 * 1024):.2f} MB")

    start_time = time.time()

    try:
        # Step 1: Detect tables
        tables = detect_tables(str(test_file))

        if not tables:
            print("\n❌ No tables detected in file")
            return

        # Find first large table
        large_table = None
        for table in tables:
            num_rows = table['end_row'] - table['start_row'] + 1
            if num_rows > MAX_ROWS_FOR_DOCLING:
                large_table = table
                break

        if not large_table:
            print(f"\n⚠️  No tables larger than {MAX_ROWS_FOR_DOCLING} rows found")
            print(f"   Using first table for testing anyway...")
            large_table = tables[0]

        print(f"\n🎯 Selected table for testing:")
        print(f"   Sheet: {large_table['sheet']}")
        print(f"   Range: {large_table['range']}")

        # Step 2: Extract header rows
        header_file = extract_header_rows(str(test_file), large_table)

        # Step 3: Convert to HTML (for validation/debugging)
        html = convert_headers_to_html(header_file)

        if not html:
            print("\n❌ HTML conversion failed")
            return

        # Step 4: Extract normalized headers directly from Excel
        headers = extract_normalized_headers(header_file, large_table)

        # Step 5: Process data rows
        data = process_data_rows(str(test_file), large_table, headers)

        # Step 6: Export results
        output_path = export_results(large_table, headers, data)

        # Step 7: Summary
        print_step(7, 7, "Test Summary")

        end_time = time.time()
        duration = end_time - start_time

        print(f"\n✅ Test completed successfully!")
        print(f"\n📊 Statistics:")
        print(f"   • Duration: {duration:.2f}s")
        print(f"   • Tables detected: {len(tables)}")
        print(f"   • Data rows processed: {len(data):,}")
        print(f"   • Output file: {output_path}")

        print(f"\n💡 Next Steps:")
        print(f"   1. Implement HTML → Worksheet parser for header extraction")
        print(f"   2. Integrate HeaderExtractor with parsed HTML structure")
        print(f"   3. Add proper column mapping with normalized headers")
        print(f"   4. Implement full data row iteration (currently limited for testing)")
        print(f"   5. Add markdown table formatting for better output visualization")

        # Cleanup temp file
        try:
            Path(header_file).unlink()
            print(f"\n🗑️  Cleaned up temporary file: {header_file}")
        except:
            pass

        print("\n" + "=" * 80)

    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()


def main():
    """Main entry point."""
    # Default test file (the large JP1 file)
    default_file = "CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.2 JP1шинхоЪщаЕчЫо.xlsx"

    if len(sys.argv) > 1:
        test_file = sys.argv[1]
    else:
        test_file = default_file
        print(f"ℹ️  No file specified, using default: {test_file}")
        print(f"   Usage: python {sys.argv[0]} [file_path]\n")

    try:
        test_large_table_strategy(test_file)
    except KeyboardInterrupt:
        print("\n\n⚠️  Test interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()