"""
Unit tests for LargeTableExtractor.

Tests the large table extraction functionality including:
- Table size classification
- Header extraction and normalization
- Chunked data processing
- Pipeline record conversion
"""

import pytest
from unittest.mock import Mock, patch, MagicMock
from uuid import uuid4

from src.extraction_v2.large_table_extractor import (
    LargeTableExtractor,
    LargeTableConfig,
    LargeTableResult,
    extract_large_tables
)


class TestLargeTableConfig:
    """Test LargeTableConfig dataclass."""

    def test_default_values(self):
        """Test default configuration values."""
        config = LargeTableConfig()

        assert config.max_rows_for_docling == 1000
        assert config.header_sample_rows == 10
        assert config.chunk_size == 100
        assert config.max_rows_to_process == -1

    def test_custom_values(self):
        """Test custom configuration values."""
        config = LargeTableConfig(
            max_rows_for_docling=500,
            header_sample_rows=5,
            chunk_size=50,
            max_rows_to_process=1000
        )

        assert config.max_rows_for_docling == 500
        assert config.header_sample_rows == 5
        assert config.chunk_size == 50
        assert config.max_rows_to_process == 1000


class TestLargeTableResult:
    """Test LargeTableResult dataclass."""

    def test_success_result(self):
        """Test successful result creation."""
        result = LargeTableResult(
            success=True,
            sheet_name="Sheet1",
            table_range="A1:Z100",
            total_rows=100,
            rows_processed=95,
            column_count=26,
            header_depth=3,
            headers={"A": "Col1", "B": "Col2"},
            records=[{"Col1": "val1"}],
            processing_time=5.5
        )

        assert result.success is True
        assert result.sheet_name == "Sheet1"
        assert result.error is None

    def test_failure_result(self):
        """Test failure result creation."""
        result = LargeTableResult(
            success=False,
            sheet_name="Sheet1",
            table_range="A1:Z100",
            total_rows=100,
            rows_processed=0,
            column_count=0,
            header_depth=0,
            headers={},
            records=[],
            error="Test error"
        )

        assert result.success is False
        assert result.error == "Test error"
        assert result.records == []


class TestLargeTableExtractor:
    """Test LargeTableExtractor class."""

    def test_init_default_config(self):
        """Test initialization with default config."""
        extractor = LargeTableExtractor()

        assert extractor.config.max_rows_for_docling == 1000
        assert extractor.border_detector is not None
        assert extractor.header_detector is not None

    def test_init_custom_config(self):
        """Test initialization with custom config."""
        config = LargeTableConfig(max_rows_for_docling=500)
        extractor = LargeTableExtractor(config)

        assert extractor.config.max_rows_for_docling == 500

    def test_is_large_table_true(self):
        """Test large table detection - table exceeds threshold."""
        extractor = LargeTableExtractor()
        table_info = {
            'start_row': 1,
            'end_row': 2000  # 2000 rows > 1000 threshold
        }

        assert extractor.is_large_table(table_info) is True

    def test_is_large_table_false(self):
        """Test large table detection - table within threshold."""
        extractor = LargeTableExtractor()
        table_info = {
            'start_row': 1,
            'end_row': 500  # 500 rows < 1000 threshold
        }

        assert extractor.is_large_table(table_info) is False

    def test_is_large_table_boundary(self):
        """Test large table detection at boundary."""
        extractor = LargeTableExtractor()

        # Exactly 1000 rows - should NOT be large
        table_info = {'start_row': 1, 'end_row': 1000}
        assert extractor.is_large_table(table_info) is False

        # 1001 rows - should be large
        table_info = {'start_row': 1, 'end_row': 1001}
        assert extractor.is_large_table(table_info) is True

    def test_to_pipeline_records(self):
        """Test conversion to pipeline record format."""
        extractor = LargeTableExtractor()
        document_id = uuid4()

        result = LargeTableResult(
            success=True,
            sheet_name="TestSheet",
            table_range="A1:C10",
            total_rows=10,
            rows_processed=8,
            column_count=3,
            header_depth=2,
            headers={"A": "Name", "B": "Value", "C": "Count"},
            records=[
                {"Name": "Row1", "Value": 100, "Count": 5},
                {"Name": "Row2", "Value": 200, "Count": 10}
            ]
        )

        pipeline_records = extractor.to_pipeline_records(result, document_id)

        assert len(pipeline_records) == 2
        assert pipeline_records[0]['document_id'] == str(document_id)
        assert pipeline_records[0]['sheet_name'] == "TestSheet"
        assert pipeline_records[0]['table_range'] == "A1:C10"
        assert pipeline_records[0]['row_index'] == 0
        assert pipeline_records[0]['data'] == {"Name": "Row1", "Value": 100, "Count": 5}
        assert pipeline_records[0]['extraction_method'] == 'large_table_chunked'
        assert pipeline_records[0]['header_depth'] == 2

        assert pipeline_records[1]['row_index'] == 1


class TestHeaderNormalization:
    """Test header normalization logic."""

    def test_remove_consecutive_duplicates(self):
        """Test that consecutive duplicate levels are removed."""
        # This tests the logic that should be in _extract_normalized_headers
        levels = ["Header", "Header", "Header"]

        # Remove consecutive duplicates
        deduplicated = []
        for level in levels:
            if not deduplicated or level != deduplicated[-1]:
                deduplicated.append(level)

        assert deduplicated == ["Header"]

    def test_remove_table_name_prefix(self):
        """Test that table name prefix is removed."""
        # Simulate multi-level header with table name
        levels = ["Table 1.2", "Section", "Column"]

        # Remove first level if multiple levels
        if len(levels) > 1:
            levels = levels[1:]

        assert levels == ["Section", "Column"]

    def test_preserve_single_level(self):
        """Test that single level headers are preserved."""
        levels = ["OnlyColumn"]

        # Should not remove anything if only one level
        if len(levels) > 1:
            levels = levels[1:]

        assert levels == ["OnlyColumn"]

    def test_join_levels(self):
        """Test that levels are joined with separator."""
        levels = ["Section", "SubSection", "Column"]
        header_name = " > ".join(levels)

        assert header_name == "Section > SubSection > Column"


class TestBorderDetectorHelpers:
    """Test BorderTableDetector helper methods."""

    def test_is_large_table(self):
        """Test BorderTableDetector.is_large_table method."""
        from src.extraction_v2.detect_border import BorderTableDetector

        detector = BorderTableDetector()

        large_table = {'start_row': 1, 'end_row': 2000}
        assert detector.is_large_table(large_table) is True

        small_table = {'start_row': 1, 'end_row': 500}
        assert detector.is_large_table(small_table) is False

    def test_classify_tables(self):
        """Test BorderTableDetector.classify_tables method."""
        from src.extraction_v2.detect_border import BorderTableDetector

        detector = BorderTableDetector()

        tables = [
            {'start_row': 1, 'end_row': 500, 'sheet': 'Small1'},
            {'start_row': 1, 'end_row': 2000, 'sheet': 'Large1'},
            {'start_row': 1, 'end_row': 300, 'sheet': 'Small2'},
            {'start_row': 1, 'end_row': 5000, 'sheet': 'Large2'},
        ]

        large, normal = detector.classify_tables(tables)

        assert len(large) == 2
        assert len(normal) == 2
        assert large[0]['sheet'] == 'Large1'
        assert large[1]['sheet'] == 'Large2'
        assert normal[0]['sheet'] == 'Small1'
        assert normal[1]['sheet'] == 'Small2'

    def test_custom_threshold(self):
        """Test BorderTableDetector with custom threshold."""
        from src.extraction_v2.detect_border import BorderTableDetector

        detector = BorderTableDetector(large_table_threshold=500)

        # 600 rows should be large with 500 threshold
        table = {'start_row': 1, 'end_row': 600}
        assert detector.is_large_table(table) is True

        # 400 rows should be normal
        table = {'start_row': 1, 'end_row': 400}
        assert detector.is_large_table(table) is False


class TestExtractLargeTablesFunction:
    """Test the extract_large_tables convenience function."""

    @patch('src.extraction_v2.large_table_extractor.LargeTableExtractor')
    def test_calls_extractor(self, mock_extractor_class):
        """Test that convenience function creates extractor and calls method."""
        mock_extractor = Mock()
        mock_extractor.extract_all_large_tables.return_value = []
        mock_extractor_class.return_value = mock_extractor

        result = extract_large_tables("test.xlsx")

        mock_extractor_class.assert_called_once_with(None)
        mock_extractor.extract_all_large_tables.assert_called_once_with(
            "test.xlsx", None
        )

    @patch('src.extraction_v2.large_table_extractor.LargeTableExtractor')
    def test_passes_config(self, mock_extractor_class):
        """Test that custom config is passed to extractor."""
        mock_extractor = Mock()
        mock_extractor.extract_all_large_tables.return_value = []
        mock_extractor_class.return_value = mock_extractor

        config = LargeTableConfig(max_rows_for_docling=500)
        result = extract_large_tables("test.xlsx", config=config)

        mock_extractor_class.assert_called_once_with(config)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])