#!/usr/bin/env python3
"""
Quick test to verify the header detection fix.
"""

import sys
from pathlib import Path

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

from src.extraction_v2.llm_header_detector import LlmHeaderDetector

# Sample problematic HTML (like table 8)
PROBLEMATIC_HTML = """
<table>
  <tr>
    <th colspan="3">上記以外</th>
    <th colspan="3" rowspan="4">－</th>
    <th colspan="6" rowspan="5">None</th>
    <th colspan="7" rowspan="6">None</th>
    <th colspan="4" rowspan="7">None</th>
    <th colspan="3">'否'</th>
  </tr>
  <tr>
    <td colspan="3" rowspan="3">－</td>
    <td colspan="3">'否'</td>
  </tr>
  <tr>
    <td colspan="3">'否'</td>
  </tr>
</table>
"""

# Good HTML (like table 4)
GOOD_HTML = """
<table>
  <tr>
    <th>ファイル番号</th>
    <th>ファイル和名</th>
    <th>ルートタグ名</th>
  </tr>
  <tr>
    <td>101</td>
    <td>事前照会（光心線提供検討依頼情報）</td>
    <td>S0741_LineEquipmentResearchObjectOpticalEquipmentConstructExamineRequest_OUT_1</td>
  </tr>
</table>
"""


def test_header_detection():
    """Test header detection on both good and problematic tables."""
    detector = LlmHeaderDetector()

    print("="*70)
    print("TEST 1: Problematic Table (data values in <th> tags)")
    print("="*70)

    # First, let's see what preprocessing does
    from bs4 import BeautifulSoup
    soup = BeautifulSoup(PROBLEMATIC_HTML, 'lxml')
    rows_before = soup.find_all('tr')
    print(f"\nBefore preprocessing:")
    print(f"  Row 0 has {len(rows_before[0].find_all('th'))} <th> tags")
    print(f"  <th> contents: {[th.get_text(strip=True) for th in rows_before[0].find_all('th')]}")

    # Apply preprocessing
    rows_after = detector._fix_misclassified_header_tags(rows_before)
    print(f"\nAfter preprocessing:")
    print(f"  Row 0 has {len(rows_after[0].find_all('th'))} <th> tags")
    print(f"  Row 0 has {len(rows_after[0].find_all('td'))} <td> tags")
    if rows_after[0].find_all('td'):
        print(f"  <td> contents: {[td.get_text(strip=True) for td in rows_after[0].find_all('td')]}")

    # Now test full header detection (with LLM)
    print(f"\nCalling LLM for header detection...")
    try:
        headers = detector.detect_headers(PROBLEMATIC_HTML)
        print(f"✓ Detected {len(headers.headers)} headers:")
        for col_idx, header in enumerate(headers.get_flat_headers()):
            print(f"    Column {col_idx}: '{header}'")
    except Exception as e:
        print(f"✗ Error: {e}")

    print("\n" + "="*70)
    print("TEST 2: Good Table (proper headers)")
    print("="*70)

    soup2 = BeautifulSoup(GOOD_HTML, 'lxml')
    rows2_before = soup2.find_all('tr')
    print(f"\nBefore preprocessing:")
    print(f"  Row 0 has {len(rows2_before[0].find_all('th'))} <th> tags")
    print(f"  <th> contents: {[th.get_text(strip=True) for th in rows2_before[0].find_all('th')]}")

    rows2_after = detector._fix_misclassified_header_tags(rows2_before)
    print(f"\nAfter preprocessing:")
    print(f"  Row 0 has {len(rows2_after[0].find_all('th'))} <th> tags")
    print(f"  Row 0 has {len(rows2_after[0].find_all('td'))} <td> tags")
    if rows2_after[0].find_all('th'):
        print(f"  <th> contents (should be unchanged): {[th.get_text(strip=True) for th in rows2_after[0].find_all('th')]}")

    print(f"\nCalling LLM for header detection...")
    try:
        headers2 = detector.detect_headers(GOOD_HTML)
        print(f"✓ Detected {len(headers2.headers)} headers:")
        for col_idx, header in enumerate(headers2.get_flat_headers()):
            print(f"    Column {col_idx}: '{header}'")
    except Exception as e:
        print(f"✗ Error: {e}")

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


if __name__ == "__main__":
    test_header_detection()