#!/usr/bin/env python3
"""Test all three types of cell merging: colspan, rowspan, and both combined."""

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_all_merge_types():
    """Test all three merge scenarios."""

    print("=" * 80)
    print("COMPREHENSIVE MERGE TEST - All Three Cases")
    print("=" * 80)

    # Test table with all three merge types
    html = """
    <table>
        <tr>
            <th>Header</th>
            <th colspan="3">Colspan Only (3 cols)</th>
            <th>Normal</th>
        </tr>
        <tr>
            <td>Row 1</td>
            <td>A</td>
            <td>B</td>
            <td>C</td>
            <td>D</td>
        </tr>
        <tr>
            <td rowspan="3">Rowspan Only (3 rows)</td>
            <td>E</td>
            <td>F</td>
            <td>G</td>
            <td>H</td>
        </tr>
        <tr>
            <td>I</td>
            <td>J</td>
            <td>K</td>
            <td>L</td>
        </tr>
        <tr>
            <td>M</td>
            <td>N</td>
            <td>O</td>
            <td>P</td>
        </tr>
        <tr>
            <td>Row 5</td>
            <td rowspan="2" colspan="3">BOTH: 2 rows × 3 cols</td>
            <td>Q</td>
        </tr>
        <tr>
            <td>Row 6</td>
            <td>R</td>
        </tr>
    </table>
    """

    print("\n📊 Test Cases:")
    print("  1. Colspan only (1 row × 3 cols)")
    print("  2. Rowspan only (3 rows × 1 col)")
    print("  3. Both combined (2 rows × 3 cols)")

    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. Analysis of merged cells:")

    # Case 1: Colspan only (row 0)
    print("\n   ✓ Case 1: Colspan Only (Header row)")
    print(f"     Row 0, Col 1-3: {table.rows[0][1:4]}")
    print(f"     Expected: ['Colspan Only (3 cols)', '', '']")
    assert table.rows[0][1] == "Colspan Only (3 cols)"
    assert table.rows[0][2] == ""
    assert table.rows[0][3] == ""
    print("     ✓ Colspan handled correctly!")

    # Case 2: Rowspan only (rows 2-4, col 0)
    print("\n   ✓ Case 2: Rowspan Only (Rows 2-4, Col 0)")
    print(f"     Row 2, Col 0: '{table.rows[2][0]}'")
    print(f"     Row 3, Col 0: '{table.rows[3][0]}'")
    print(f"     Row 4, Col 0: '{table.rows[4][0]}'")
    print(f"     Expected: All should be 'Rowspan Only (3 rows)'")
    assert table.rows[2][0] == "Rowspan Only (3 rows)"
    assert table.rows[3][0] == "Rowspan Only (3 rows)"
    assert table.rows[4][0] == "Rowspan Only (3 rows)"
    print("     ✓ Rowspan handled correctly!")

    # Case 3: Both rowspan and colspan (rows 5-6, cols 1-3)
    print("\n   ✓ Case 3: Both Rowspan AND Colspan (2 rows × 3 cols)")
    print(f"     Row 5, Cols 1-3: {table.rows[5][1:4]}")
    print(f"     Row 6, Cols 1-3: {table.rows[6][1:4]}")
    print(f"     Expected:")
    print(f"       Row 5: ['BOTH: 2 rows × 3 cols', '', '']")
    print(f"       Row 6: ['BOTH: 2 rows × 3 cols', '', '']")

    assert table.rows[5][1] == "BOTH: 2 rows × 3 cols"
    assert table.rows[5][2] == ""
    assert table.rows[5][3] == ""
    assert table.rows[6][1] == "BOTH: 2 rows × 3 cols"
    assert table.rows[6][2] == ""
    assert table.rows[6][3] == ""
    print("     ✓ Combined rowspan+colspan handled correctly!")

    print("\n4. Creating mock headers...")
    headers = HeaderStructure(
        header_rows=[0],
        headers=[
            {"column": 0, "hierarchy": ["Header"]},
            {"column": 1, "hierarchy": ["Colspan Only (3 cols)"]},
            {"column": 2, "hierarchy": ["Colspan Only (3 cols)"]},
            {"column": 3, "hierarchy": ["Colspan Only (3 cols)"]},
            {"column": 4, "hierarchy": ["Normal"]},
        ]
    )

    print("\n5. 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("\n6. Verifying records have no empty strings:")
    for i, record in enumerate(records):
        empty_count = sum(1 for v in record['content'].values() if v == '')
        non_empty_count = len(record['content'])

        print(f"\n   Record {i + 1}:")
        print(f"     Keys: {list(record['content'].keys())}")
        print(f"     Values: {list(record['content'].values())}")

        if empty_count > 0:
            print(f"     ⚠️  WARNING: {empty_count} empty strings found!")
        else:
            print(f"     ✓ Clean: {non_empty_count} non-empty values")

    print("\n" + "=" * 80)
    print("TEST COMPLETE - ALL THREE MERGE TYPES VERIFIED!")
    print("=" * 80)

    # Final summary
    total_empty = sum(
        sum(1 for v in record['content'].values() if v == '')
        for record in records
    )

    print(f"\nFinal Summary:")
    print(f"  ✓ Colspan only: Handled")
    print(f"  ✓ Rowspan only: Handled")
    print(f"  ✓ Both combined: Handled")
    print(f"  ✓ Empty strings in records: {total_empty}")

    if total_empty == 0:
        print(f"\n✓✓✓ SUCCESS: All merge types work perfectly with no empty strings!")
    else:
        print(f"\n⚠️  WARNING: Found {total_empty} empty strings")


if __name__ == "__main__":
    test_all_merge_types()
