#!/usr/bin/env python3
"""Test merged cell normalization in V2 extraction pipeline."""

from bs4 import BeautifulSoup
from uuid import uuid4

from src.extraction_v2.html_table_extractor import HtmlTableExtractor
from src.extraction_v2.llm_header_detector import HeaderStructure
from src.extraction_v2.record_builder import RecordBuilder


def test_merged_cell_normalization():
    """Test that merged cells are properly normalized in records."""

    print("=" * 70)
    print("MERGED CELL NORMALIZATION TEST")
    print("=" * 70)

    # Create test HTML with merged cells
    html = """
    <table>
        <tr>
            <th>Item</th>
            <th colspan="3">Specifications</th>
            <th>Price</th>
        </tr>
        <tr>
            <th>Name</th>
            <th>Width</th>
            <th>Height</th>
            <th>Depth</th>
            <th>USD</th>
        </tr>
        <tr>
            <td rowspan="2">Widget</td>
            <td>10</td>
            <td>20</td>
            <td>30</td>
            <td>100</td>
        </tr>
        <tr>
            <td>15</td>
            <td>25</td>
            <td>35</td>
            <td>150</td>
        </tr>
        <tr>
            <td colspan="4">Special Item with long description</td>
            <td>200</td>
        </tr>
    </table>
    """

    print("\n1. Extracting tables from HTML...")
    extractor = HtmlTableExtractor()
    tables = extractor.extract_tables(html)

    if not tables:
        print("ERROR: No tables found")
        return

    table = tables[0]
    print(f"   Table extracted: {table.row_count} rows x {table.col_count} cols")

    print("\n2. Raw 2D grid after merged cell expansion:")
    for i, row in enumerate(table.rows):
        print(f"   Row {i}: {row}")

    print("\n3. Creating mock headers...")
    # Create simple headers for testing
    headers = HeaderStructure(
        header_rows=[0, 1],
        headers=[
            {"column": 0, "hierarchy": ["Item", "Name"]},
            {"column": 1, "hierarchy": ["Specifications", "Width"]},
            {"column": 2, "hierarchy": ["Specifications", "Height"]},
            {"column": 3, "hierarchy": ["Specifications", "Depth"]},
            {"column": 4, "hierarchy": ["Price", "USD"]},
        ]
    )

    flat_headers = headers.get_flat_headers()
    print(f"   Flat headers: {flat_headers}")

    print("\n4. Building records with merged cell normalization...")
    builder = RecordBuilder()
    records = builder.build_records(
        table=table,
        headers=headers,
        document_id=uuid4(),
        sheet_name="Test"
    )

    print(f"   Built {len(records)} records")

    print("\n5. Examining records (empty values should be skipped):")
    for i, record in enumerate(records):
        print(f"\n   Record {i + 1}:")
        print(f"   Content: {record['content']}")
        print(f"   Source: {record['_source']}")

        # Verify no empty values
        empty_count = sum(1 for v in record['content'].values() if v == '')
        if empty_count > 0:
            print(f"   ⚠️  WARNING: Found {empty_count} empty string values!")
        else:
            print(f"   ✓ Clean: No empty string values")

    print("\n" + "=" * 70)
    print("TEST COMPLETE")
    print("=" * 70)

    # Assertions
    assert len(records) == 3, f"Expected 3 records, got {len(records)}"

    # Check first record (Widget with rowspan)
    record1 = records[0]
    assert 'Item -> Name' in record1['content']
    assert record1['content']['Item -> Name'] == 'Widget'
    assert '' not in record1['content'].values(), "Record should not contain empty strings"

    # Check third record (colspan for description)
    record3 = records[2]
    assert 'Item -> Name' in record3['content']
    assert record3['content']['Item -> Name'] == 'Special Item with long description'
    assert '' not in record3['content'].values(), "Record should not contain empty strings"

    print("\n✓ All assertions passed!")


if __name__ == "__main__":
    test_merged_cell_normalization()