"""
Unit tests for ParseDocumentOperator.
"""

import pytest
from unittest.mock import Mock, MagicMock, patch
from plugins.operators.parse_document import ParseDocumentOperator


class TestParseDocumentOperator:
    """Test suite for ParseDocumentOperator."""

    @pytest.fixture
    def mock_context(self):
        """Create a mock Airflow context."""
        context = {
            'ti': Mock(),
        }
        return context

    @pytest.fixture
    def sample_fetch_result(self):
        """Sample fetch result data."""
        return {
            'local_path': '/tmp/test/sample.xlsx',
            'content_type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            'original_filename': 'sample.xlsx',
            'file_extension': 'xlsx',
            'document_id': 'test-doc-123',
        }

    def test_parse_excel_document_normal_tables(self, mock_context, sample_fetch_result):
        """Test parsing Excel document with normal tables (no large tables)."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Mock large table extractor - returns normal tables only
            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = [
                {'sheet': 'Sheet1', 'is_large': False, 'row_count': 50}
            ]
            mock_large_extractor_class.return_value = mock_large_extractor

            # Mock converter
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '<html><table>...</table></html>'
            mock_converter_class.return_value = mock_converter

            # Mock table extractor
            mock_table1 = MagicMock()
            mock_table1.html = '<table>Table 1</table>'
            mock_table1.rows = [['A', 'B'], ['1', '2']]

            mock_raw_table = MagicMock()
            mock_raw_table.rows = [['Raw', 'Data']]

            mock_extractor = MagicMock()
            mock_extractor.extract_all_tables.return_value = ([mock_table1], [mock_raw_table])
            mock_extractor_class.return_value = mock_extractor

            op = ParseDocumentOperator(
                task_id='test_parse',
                doc_type='excel',
                fetch_document_task_id='fetch_document'
            )

            result = op.execute(mock_context)

            # Verify XCom pull
            mock_context['ti'].xcom_pull.assert_called_once_with(task_ids='fetch_document')

            # Verify large table detection was called
            mock_large_extractor.detect_tables.assert_called_once_with('/tmp/test/sample.xlsx')

            # Verify converter was used (for normal tables)
            mock_converter.convert_to_html.assert_called_once_with('/tmp/test/sample.xlsx')

            # Verify extractor was used
            mock_extractor.extract_all_tables.assert_called_once()

            # Verify result structure
            assert result['doc_type'] == 'excel'
            assert result['html_content'] == '<html><table>...</table></html>'
            assert result['markdown_content'] is None
            assert len(result['tables']) == 1
            assert result['tables'][0]['html'] == '<table>Table 1</table>'
            assert len(result['raw_text_tables']) == 1
            assert result['large_table_records'] == []
            assert result['document_id'] == 'test-doc-123'
            assert result['local_path'] == '/tmp/test/sample.xlsx'

    def test_parse_excel_document_with_large_tables(self, mock_context, sample_fetch_result):
        """Test parsing Excel document with large tables (>1000 rows)."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Mock large table extractor - returns large tables only
            mock_large_result = MagicMock()
            mock_large_result.success = True
            mock_large_result.records = [
                {'row': 1, 'data': 'large_record_1'},
                {'row': 2, 'data': 'large_record_2'},
            ]
            mock_large_result.rows_processed = 1500
            mock_large_result.sheet_name = 'LargeSheet'

            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = [
                {'sheet': 'LargeSheet', 'is_large': True, 'row_count': 1500}
            ]
            mock_large_extractor.extract_table.return_value = mock_large_result
            mock_large_extractor_class.return_value = mock_large_extractor

            # Converter and extractor should NOT be called (no normal tables)
            mock_converter = MagicMock()
            mock_converter_class.return_value = mock_converter

            mock_extractor = MagicMock()
            mock_extractor_class.return_value = mock_extractor

            op = ParseDocumentOperator(
                task_id='test_parse',
                doc_type='excel',
            )

            result = op.execute(mock_context)

            # Verify large table detection was called
            mock_large_extractor.detect_tables.assert_called_once_with('/tmp/test/sample.xlsx')

            # Verify large table extraction was called
            mock_large_extractor.extract_table.assert_called_once()

            # Verify converter was NOT called (only large tables)
            mock_converter.convert_to_html.assert_not_called()

            # Verify result structure
            assert result['doc_type'] == 'excel'
            assert result['html_content'] is None  # No Docling conversion for large-only
            assert result['tables'] == []
            assert result['raw_text_tables'] == []
            assert len(result['large_table_records']) == 2
            assert result['large_table_records'][0]['data'] == 'large_record_1'

    def test_parse_excel_document_mixed_tables(self, mock_context, sample_fetch_result):
        """Test parsing Excel document with both large and normal tables."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Mock large table extractor - returns both large and normal
            mock_large_result = MagicMock()
            mock_large_result.success = True
            mock_large_result.records = [{'data': 'large_record'}]
            mock_large_result.rows_processed = 2000
            mock_large_result.sheet_name = 'LargeSheet'

            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = [
                {'sheet': 'LargeSheet', 'is_large': True, 'row_count': 2000},
                {'sheet': 'NormalSheet', 'is_large': False, 'row_count': 100},
            ]
            mock_large_extractor.extract_table.return_value = mock_large_result
            mock_large_extractor_class.return_value = mock_large_extractor

            # Mock converter for normal tables
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '<html><table>normal</table></html>'
            mock_converter_class.return_value = mock_converter

            # Mock table extractor
            mock_table = MagicMock()
            mock_table.html = '<table>Normal Table</table>'
            mock_table.rows = [['A'], ['1']]

            mock_extractor = MagicMock()
            mock_extractor.extract_all_tables.return_value = ([mock_table], [])
            mock_extractor_class.return_value = mock_extractor

            op = ParseDocumentOperator(task_id='test_parse', doc_type='excel')
            result = op.execute(mock_context)

            # Verify both paths were executed
            mock_large_extractor.detect_tables.assert_called_once()
            mock_large_extractor.extract_table.assert_called_once()
            mock_converter.convert_to_html.assert_called_once()
            mock_extractor.extract_all_tables.assert_called_once()

            # Verify result includes both
            assert len(result['large_table_records']) == 1
            assert len(result['tables']) == 1
            assert result['html_content'] == '<html><table>normal</table></html>'

    def test_parse_excel_large_table_extraction_failure(self, mock_context, sample_fetch_result):
        """Test handling of large table extraction failure."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Mock large table extractor - extraction fails
            mock_large_result = MagicMock()
            mock_large_result.success = False
            mock_large_result.error = 'Extraction failed'

            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = [
                {'sheet': 'LargeSheet', 'is_large': True, 'row_count': 1500}
            ]
            mock_large_extractor.extract_table.return_value = mock_large_result
            mock_large_extractor_class.return_value = mock_large_extractor

            mock_converter = MagicMock()
            mock_converter_class.return_value = mock_converter

            mock_extractor = MagicMock()
            mock_extractor_class.return_value = mock_extractor

            op = ParseDocumentOperator(task_id='test_parse', doc_type='excel')
            result = op.execute(mock_context)

            # Verify extraction was attempted
            mock_large_extractor.extract_table.assert_called_once()

            # Result should have empty large_table_records due to failure
            assert result['large_table_records'] == []

    def test_parse_excel_no_tables_detected(self, mock_context, sample_fetch_result):
        """Test Excel parsing when border detection finds no tables."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Mock large table extractor - no tables detected
            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = []
            mock_large_extractor_class.return_value = mock_large_extractor

            # Should fall back to Docling conversion
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '<html>fallback</html>'
            mock_converter_class.return_value = mock_converter

            mock_extractor = MagicMock()
            mock_extractor.extract_all_tables.return_value = ([], [])
            mock_extractor_class.return_value = mock_extractor

            op = ParseDocumentOperator(task_id='test_parse', doc_type='excel')
            result = op.execute(mock_context)

            # Verify fallback to Docling was used
            mock_converter.convert_to_html.assert_called_once()
            assert result['html_content'] == '<html>fallback</html>'

    def test_parse_word_document(self, mock_context, sample_fetch_result):
        """Test parsing Word document to Markdown."""
        fetch_result = sample_fetch_result.copy()
        fetch_result['file_extension'] = 'docx'
        fetch_result['local_path'] = '/tmp/test/document.docx'

        mock_context['ti'].xcom_pull.return_value = fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class:
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '# Document Title\n\nContent here...'
            mock_converter_class.return_value = mock_converter

            op = ParseDocumentOperator(
                task_id='test_parse',
                doc_type='word'
            )

            result = op.execute(mock_context)

            # Verify converter was used
            mock_converter.convert_to_html.assert_called_once_with('/tmp/test/document.docx')

            # Verify result structure
            assert result['doc_type'] == 'word'
            assert result['html_content'] is None
            assert result['markdown_content'] == '# Document Title\n\nContent here...'
            assert result['tables'] == []
            assert result['raw_text_tables'] == []
            assert result['document_id'] == 'test-doc-123'

    def test_parse_pdf_document(self, mock_context, sample_fetch_result):
        """Test parsing PDF document to Markdown."""
        fetch_result = sample_fetch_result.copy()
        fetch_result['file_extension'] = 'pdf'
        fetch_result['local_path'] = '/tmp/test/document.pdf'

        mock_context['ti'].xcom_pull.return_value = fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class:
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '# PDF Content\n\nParsed text'
            mock_converter_class.return_value = mock_converter

            op = ParseDocumentOperator(
                task_id='test_parse',
                doc_type='pdf'
            )

            result = op.execute(mock_context)

            assert result['doc_type'] == 'pdf'
            assert result['markdown_content'] == '# PDF Content\n\nParsed text'

    def test_parse_powerpoint_document(self, mock_context, sample_fetch_result):
        """Test parsing PowerPoint document to Markdown."""
        fetch_result = sample_fetch_result.copy()
        fetch_result['file_extension'] = 'pptx'
        fetch_result['local_path'] = '/tmp/test/presentation.pptx'

        mock_context['ti'].xcom_pull.return_value = fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class:
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '# Slide 1\n\nContent'
            mock_converter_class.return_value = mock_converter

            op = ParseDocumentOperator(
                task_id='test_parse',
                doc_type='powerpoint'
            )

            result = op.execute(mock_context)

            assert result['doc_type'] == 'powerpoint'
            assert result['markdown_content'] == '# Slide 1\n\nContent'

    def test_parse_no_fetch_result_raises(self, mock_context):
        """Test that missing fetch result raises ValueError."""
        mock_context['ti'].xcom_pull.return_value = None

        op = ParseDocumentOperator(task_id='test_parse', doc_type='excel')

        with pytest.raises(ValueError, match='No fetch result found'):
            op.execute(mock_context)

    def test_parse_conversion_failure_raises(self, mock_context, sample_fetch_result):
        """Test that conversion failure raises exception."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Mock large table extractor - returns normal table to trigger Docling
            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = [
                {'sheet': 'Sheet1', 'is_large': False}
            ]
            mock_large_extractor_class.return_value = mock_large_extractor

            mock_converter = MagicMock()
            mock_converter.convert_to_html.side_effect = Exception('Conversion failed')
            mock_converter_class.return_value = mock_converter

            op = ParseDocumentOperator(task_id='test_parse', doc_type='excel')

            with pytest.raises(Exception, match='Conversion failed'):
                op.execute(mock_context)

    def test_custom_fetch_task_id(self, mock_context, sample_fetch_result):
        """Test using custom fetch task ID."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class:
            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = 'Content'
            mock_converter_class.return_value = mock_converter

            op = ParseDocumentOperator(
                task_id='test_parse',
                doc_type='word',
                fetch_document_task_id='custom_fetch'
            )

            result = op.execute(mock_context)

            # Verify custom task_id was used
            mock_context['ti'].xcom_pull.assert_called_once_with(task_ids='custom_fetch')
            assert result['doc_type'] == 'word'

    def test_excel_with_no_tables(self, mock_context, sample_fetch_result):
        """Test Excel parsing when no tables are found (after Docling conversion)."""
        mock_context['ti'].xcom_pull.return_value = sample_fetch_result

        with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \
             patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \
             patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class:

            # Large table extractor finds normal tables
            mock_large_extractor = MagicMock()
            mock_large_extractor.detect_tables.return_value = [
                {'sheet': 'Sheet1', 'is_large': False}
            ]
            mock_large_extractor_class.return_value = mock_large_extractor

            mock_converter = MagicMock()
            mock_converter.convert_to_html.return_value = '<html>No tables</html>'
            mock_converter_class.return_value = mock_converter

            mock_extractor = MagicMock()
            mock_extractor.extract_all_tables.return_value = ([], [])  # No tables
            mock_extractor_class.return_value = mock_extractor

            op = ParseDocumentOperator(task_id='test_parse', doc_type='excel')
            result = op.execute(mock_context)

            assert result['tables'] == []
            assert result['raw_text_tables'] == []
            assert result['html_content'] == '<html>No tables</html>'
            assert result['large_table_records'] == []
