"""
Unit tests for ChunkContentOperator.
"""

import pytest
from unittest.mock import Mock, MagicMock, patch
from uuid import UUID
from plugins.operators.chunk_content import ChunkContentOperator


class TestChunkContentOperator:
    """Test suite for ChunkContentOperator."""

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

    @pytest.fixture
    def sample_validate_result(self):
        """Sample validate result data."""
        return {
            'bucket': 'documents',
            'object_key': 'test/sample.xlsx',
            'content_type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            'document_id': '550e8400-e29b-41d4-a716-446655440000',
        }

    @pytest.fixture
    def sample_excel_parse_result(self):
        """Sample Excel parse result data."""
        return {
            'doc_type': 'excel',
            'tables': [
                {
                    'html': '<table><tr><th>Name</th></tr><tr><td>Alice</td></tr></table>',
                    'rows': [['Name'], ['Alice']],
                },
            ],
            'raw_text_tables': [
                {'rows': [['Some raw text'], ['More text']]},
            ],
            'local_path': '/tmp/test/sample.xlsx',
            'document_id': '550e8400-e29b-41d4-a716-446655440000',
        }

    @pytest.fixture
    def sample_headers_result(self):
        """Sample headers detection result."""
        return {
            'tables_processed': 1,
            'header_data': [
                {
                    'table_index': 0,
                    'header_rows': [0],
                    'headers': [{'column': 0, 'hierarchy': ['Name']}],
                }
            ],
        }

    def test_chunk_excel_success(
        self,
        mock_context,
        sample_validate_result,
        sample_excel_parse_result,
        sample_headers_result
    ):
        """Test successful Excel table chunking."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            elif task_ids == 'parse_excel':
                return sample_excel_parse_result
            elif task_ids == 'detect_headers':
                return sample_headers_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.RecordBuilder') as mock_builder_class, \
             patch('plugins.operators.chunk_content.HeaderStructure') as mock_header_struct, \
             patch('plugins.operators.chunk_content.HtmlTable') as mock_html_table:

            # Mock RecordBuilder
            mock_builder = MagicMock()
            mock_builder.build_records.return_value = [
                {'content': {'Name': 'Alice'}, '_source': {'document_id': 'test', 'sheet': 'Sheet1', 'row': 1, 'table_id': 0}},
            ]
            mock_builder.build_raw_text_records.return_value = [
                {'content': {'text': 'Some raw text\nMore text'}, '_source': {'document_id': 'test', 'sheet': 'Unknown', 'row': -1, 'table_id': 1}},
            ]
            mock_builder_class.return_value = mock_builder

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='table',
            )

            result = op.execute(mock_context)

            # Verify result structure
            assert result['chunking_strategy'] == 'table'
            assert result['chunk_count'] == 2
            assert len(result['chunks']) == 2

            # Verify RecordBuilder was used
            mock_builder.build_records.assert_called_once()
            mock_builder.build_raw_text_records.assert_called_once()

    def test_chunk_excel_no_parse_result(self, mock_context, sample_validate_result):
        """Test handling of missing parse result."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.RecordBuilder'), \
             patch('plugins.operators.chunk_content.HeaderStructure'), \
             patch('plugins.operators.chunk_content.HtmlTable'):

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='table',
            )

            result = op.execute(mock_context)

            assert result['chunk_count'] == 0
            assert result['chunks'] == []
            assert result['chunking_strategy'] == 'table'

    def test_chunk_excel_no_tables(self, mock_context, sample_validate_result):
        """Test handling of Excel with no tables."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            elif task_ids == 'parse_excel':
                return {'doc_type': 'excel', 'tables': [], 'raw_text_tables': []}
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.RecordBuilder'), \
             patch('plugins.operators.chunk_content.HeaderStructure'), \
             patch('plugins.operators.chunk_content.HtmlTable'):

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='table',
            )

            result = op.execute(mock_context)

            assert result['chunk_count'] == 0
            assert result['chunks'] == []

    def test_chunk_semantic_success(self, mock_context, sample_validate_result):
        """Test successful semantic text chunking."""
        sample_word_result = {
            'doc_type': 'word',
            'markdown_content': '# Document\n\nSome content',
            'local_path': '/tmp/test/document.docx',
            'original_filename': 'document.docx',
        }

        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            elif task_ids == 'parse_word':
                return sample_word_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.SemanticChunker') as mock_chunker_class, \
             patch('docling.document_converter.DocumentConverter') as mock_converter_class:

            # Mock SemanticChunker
            mock_chunker = MagicMock()
            mock_chunker.chunk_document.return_value = [
                {'file_id': 'test', 'text': '# Document\n\nSome content', 'metadata': {}},
            ]
            mock_chunker_class.return_value = mock_chunker

            # Mock DocumentConverter
            mock_converter = MagicMock()
            mock_result = MagicMock()
            mock_result.document = MagicMock()
            mock_converter.convert.return_value = mock_result
            mock_converter_class.return_value = mock_converter

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='semantic',
            )

            result = op.execute(mock_context)

            # Verify result structure
            assert result['chunking_strategy'] == 'semantic'
            assert result['chunk_count'] == 1
            assert len(result['chunks']) == 1

    def test_chunk_semantic_no_parse_result(self, mock_context, sample_validate_result):
        """Test semantic chunking with no parse result."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.SemanticChunker'):
            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='semantic',
            )

            result = op.execute(mock_context)

            assert result['chunk_count'] == 0
            assert result['chunks'] == []
            assert result['chunking_strategy'] == 'semantic'

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

        op = ChunkContentOperator(
            task_id='test_chunk',
            chunking_strategy='table',
        )

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

    def test_custom_max_tokens(self, mock_context, sample_validate_result):
        """Test custom max_tokens parameter."""
        op = ChunkContentOperator(
            task_id='test_chunk',
            chunking_strategy='semantic',
            max_tokens=4000,
        )

        assert op.max_tokens == 4000

    def test_find_headers_for_table(self):
        """Test _find_headers_for_table helper method."""
        op = ChunkContentOperator(
            task_id='test_chunk',
            chunking_strategy='table',
        )

        header_data = [
            {'table_index': 0, 'header_rows': [0], 'headers': [{'column': 0}]},
            {'table_index': 1, 'header_rows': [0, 1], 'headers': [{'column': 0}, {'column': 1}]},
        ]

        # Find existing table
        result = op._find_headers_for_table(header_data, 1)
        assert result['table_index'] == 1
        assert result['header_rows'] == [0, 1]

        # Fallback for non-existing table
        result = op._find_headers_for_table(header_data, 99)
        assert result['table_index'] == 99
        assert result['header_rows'] == [0]
        assert result['headers'] == []

    def test_empty_result(self):
        """Test _empty_result helper method."""
        op = ChunkContentOperator(
            task_id='test_chunk',
            chunking_strategy='table',
        )

        result = op._empty_result('semantic')
        assert result['chunks'] == []
        assert result['chunk_count'] == 0
        assert result['chunking_strategy'] == 'semantic'

    def test_custom_task_ids(self, mock_context, sample_validate_result):
        """Test using custom task IDs."""
        op = ChunkContentOperator(
            task_id='test_chunk',
            chunking_strategy='table',
            validate_task_id='custom_validate',
            parse_task_id='custom_parse',
            detect_headers_task_id='custom_headers',
        )

        assert op.validate_task_id == 'custom_validate'
        assert op.parse_task_id == 'custom_parse'
        assert op.detect_headers_task_id == 'custom_headers'

    def test_default_constants(self):
        """Test default constant values."""
        assert ChunkContentOperator.DEFAULT_MAX_TOKENS == 8000
        assert ChunkContentOperator.DEFAULT_VALIDATE_TASK_ID == 'validate_event'
        assert ChunkContentOperator.DEFAULT_PARSE_EXCEL_TASK_ID == 'parse_excel'
        assert ChunkContentOperator.DEFAULT_DETECT_HEADERS_TASK_ID == 'detect_headers'
        assert ChunkContentOperator.VALID_STRATEGIES == ('table', 'semantic')

    def test_invalid_chunking_strategy_raises(self):
        """Test that invalid chunking_strategy raises ValueError."""
        with pytest.raises(ValueError, match="Invalid chunking_strategy"):
            ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='invalid_strategy',
            )

    def test_chunk_excel_with_large_table_records(
        self,
        mock_context,
        sample_validate_result,
        sample_headers_result
    ):
        """Test chunking Excel with pre-processed large table records."""
        # Parse result includes large_table_records from ParseDocumentOperator
        parse_result_with_large = {
            'doc_type': 'excel',
            'tables': [],  # No normal tables
            'raw_text_tables': [],
            'large_table_records': [
                {'data': 'large_record_1', 'row': 1},
                {'data': 'large_record_2', 'row': 2},
                {'data': 'large_record_3', 'row': 3},
            ],
            'local_path': '/tmp/test/large.xlsx',
            'document_id': '550e8400-e29b-41d4-a716-446655440000',
        }

        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            elif task_ids == 'parse_excel':
                return parse_result_with_large
            elif task_ids == 'detect_headers':
                return sample_headers_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.RecordBuilder'), \
             patch('plugins.operators.chunk_content.HeaderStructure'), \
             patch('plugins.operators.chunk_content.HtmlTable'):

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='table',
            )

            result = op.execute(mock_context)

            # Verify large table records are included
            assert result['chunk_count'] == 3
            assert len(result['chunks']) == 3
            assert result['chunks'][0]['data'] == 'large_record_1'
            assert result['chunks'][1]['data'] == 'large_record_2'
            assert result['chunks'][2]['data'] == 'large_record_3'

    def test_chunk_excel_mixed_normal_and_large_tables(
        self,
        mock_context,
        sample_validate_result,
        sample_headers_result
    ):
        """Test chunking Excel with both normal tables and large table records."""
        # Parse result includes both normal tables and large_table_records
        mixed_parse_result = {
            'doc_type': 'excel',
            'tables': [
                {
                    'html': '<table><tr><th>Name</th></tr><tr><td>Bob</td></tr></table>',
                    'rows': [['Name'], ['Bob']],
                },
            ],
            'raw_text_tables': [],
            'large_table_records': [
                {'data': 'large_record', 'row': 1000},
            ],
            'local_path': '/tmp/test/mixed.xlsx',
            'document_id': '550e8400-e29b-41d4-a716-446655440000',
        }

        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            elif task_ids == 'parse_excel':
                return mixed_parse_result
            elif task_ids == 'detect_headers':
                return sample_headers_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.RecordBuilder') as mock_builder_class, \
             patch('plugins.operators.chunk_content.HeaderStructure'), \
             patch('plugins.operators.chunk_content.HtmlTable'):

            # Mock RecordBuilder for normal table
            mock_builder = MagicMock()
            mock_builder.build_records.return_value = [
                {'content': {'Name': 'Bob'}, '_source': {'row': 1}},
            ]
            mock_builder.build_raw_text_records.return_value = []
            mock_builder_class.return_value = mock_builder

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='table',
            )

            result = op.execute(mock_context)

            # Large table records (1) + normal table records (1) = 2
            assert result['chunk_count'] == 2
            assert len(result['chunks']) == 2

            # First chunk should be large table record (added first)
            assert result['chunks'][0]['data'] == 'large_record'

            # Second chunk should be normal table record
            assert result['chunks'][1]['content']['Name'] == 'Bob'

    def test_chunk_excel_empty_large_table_records(
        self,
        mock_context,
        sample_validate_result,
        sample_headers_result
    ):
        """Test that empty large_table_records doesn't cause issues."""
        parse_result_empty_large = {
            'doc_type': 'excel',
            'tables': [],
            'raw_text_tables': [],
            'large_table_records': [],  # Empty list
            'local_path': '/tmp/test/empty.xlsx',
            'document_id': '550e8400-e29b-41d4-a716-446655440000',
        }

        def xcom_pull_side_effect(task_ids):
            if task_ids == 'validate_event':
                return sample_validate_result
            elif task_ids == 'parse_excel':
                return parse_result_empty_large
            elif task_ids == 'detect_headers':
                return sample_headers_result
            return None

        mock_context['ti'].xcom_pull.side_effect = xcom_pull_side_effect

        with patch('plugins.operators.chunk_content.RecordBuilder'), \
             patch('plugins.operators.chunk_content.HeaderStructure'), \
             patch('plugins.operators.chunk_content.HtmlTable'):

            op = ChunkContentOperator(
                task_id='test_chunk',
                chunking_strategy='table',
            )

            result = op.execute(mock_context)

            # No tables at all = empty result
            assert result['chunk_count'] == 0
            assert result['chunks'] == []