"""
Unit tests for StoreVectorsOperator.
"""

import pytest
from unittest.mock import Mock, MagicMock, patch, AsyncMock
from uuid import UUID
from plugins.operators.store_vectors import StoreVectorsOperator, ExcelRecordProxy


class TestStoreVectorsOperator:
    """Test suite for StoreVectorsOperator."""

    @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_embedding_result_excel(self):
        """Sample embedding result with Excel table chunks."""
        return {
            'embeddings': [
                {'chunk_id': 'doc_t0_r1', 'vector': [0.1] * 3072},
                {'chunk_id': 'doc_t0_r2', 'vector': [0.2] * 3072},
            ],
            'chunks': [
                {
                    'content': {'Name': 'Alice', 'Age': '25'},
                    '_source': {
                        'document_id': '550e8400-e29b-41d4-a716-446655440000',
                        'filename': 'sample.xlsx',
                        'sheet': 'Sheet1',
                        'table_id': 0,
                        'row': 1,
                        'col_range': 'A:B',
                    },
                },
                {
                    'content': {'Name': 'Bob', 'Age': '30'},
                    '_source': {
                        'document_id': '550e8400-e29b-41d4-a716-446655440000',
                        'filename': 'sample.xlsx',
                        'sheet': 'Sheet1',
                        'table_id': 0,
                        'row': 2,
                        'col_range': 'A:B',
                    },
                },
            ],
            'embedding_count': 2,
            'model': 'text-embedding-3-large',
        }

    @pytest.fixture
    def sample_embedding_result_word(self):
        """Sample embedding result with semantic chunks."""
        return {
            'embeddings': [
                {'chunk_id': '550e8400-e29b-41d4-a716-446655440001', 'vector': [0.1] * 3072},
            ],
            'chunks': [
                {
                    'text': 'This is the document content.',
                    'metadata': {'page_number': 1},
                    'original_filename': 'document.docx',
                },
            ],
            'embedding_count': 1,
            'model': 'text-embedding-3-large',
        }

    def test_store_vectors_success_excel(
        self,
        mock_context,
        sample_validate_result,
        sample_embedding_result_excel
    ):
        """Test successful vector storage for Excel documents."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'generate_embeddings':
                return sample_embedding_result_excel
            elif 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.store_vectors.VectorStore') as mock_store_class:
            # Mock vector store
            mock_store = MagicMock()
            mock_store.create_collection = AsyncMock()
            mock_store.index_with_vectors = AsyncMock(return_value=2)
            mock_store._disconnect = MagicMock()
            mock_store.collection_name = 'test_collection'
            mock_store_class.return_value = mock_store

            op = StoreVectorsOperator(task_id='test_store_vectors')

            result = op.execute(mock_context)

            # Verify result structure
            assert result['status'] == 'success'
            assert result['stored_count'] == 2
            assert result['document_id'] == '550e8400-e29b-41d4-a716-446655440000'
            assert result['document_type'] == 'excel'
            assert result['collection'] == 'test_collection'

            # Verify VectorStore methods were called
            mock_store.create_collection.assert_called_once()
            mock_store.index_with_vectors.assert_called_once()

    def test_store_vectors_success_word(
        self,
        mock_context,
        sample_validate_result,
        sample_embedding_result_word
    ):
        """Test successful vector storage for Word documents."""
        # Update validate result for Word document
        sample_validate_result['object_key'] = 'test/document.docx'

        def xcom_pull_side_effect(task_ids):
            if task_ids == 'generate_embeddings':
                return sample_embedding_result_word
            elif 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.store_vectors.VectorStore') as mock_store_class:
            mock_store = MagicMock()
            mock_store.create_collection = AsyncMock()
            mock_store.index_with_vectors = AsyncMock(return_value=1)
            mock_store._disconnect = MagicMock()
            mock_store.collection_name = 'test_collection'
            mock_store_class.return_value = mock_store

            op = StoreVectorsOperator(task_id='test_store_vectors')

            result = op.execute(mock_context)

            assert result['status'] == 'success'
            assert result['stored_count'] == 1
            assert result['document_type'] == 'word'

    def test_store_vectors_no_embeddings(self, mock_context, sample_validate_result):
        """Test handling of empty embeddings."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'generate_embeddings':
                return {'embeddings': [], 'chunks': []}
            elif task_ids == 'validate_event':
                return sample_validate_result
            return None

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

        op = StoreVectorsOperator(task_id='test_store_vectors')

        result = op.execute(mock_context)

        assert result['status'] == 'empty'
        assert result['stored_count'] == 0

    def test_store_vectors_no_embedding_result(self, mock_context, sample_validate_result):
        """Test handling of missing embedding 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

        op = StoreVectorsOperator(task_id='test_store_vectors')

        result = op.execute(mock_context)

        assert result['status'] == 'empty'
        assert result['stored_count'] == 0

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

        op = StoreVectorsOperator(task_id='test_store_vectors')

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

    def test_custom_task_ids(self):
        """Test using custom task IDs."""
        op = StoreVectorsOperator(
            task_id='test_store',
            embeddings_task_id='custom_embeddings',
            validate_task_id='custom_validate',
        )

        assert op.embeddings_task_id == 'custom_embeddings'
        assert op.validate_task_id == 'custom_validate'

    def test_default_constants(self):
        """Test default constant values."""
        assert StoreVectorsOperator.DEFAULT_EMBEDDINGS_TASK_ID == 'generate_embeddings'
        assert StoreVectorsOperator.DEFAULT_VALIDATE_TASK_ID == 'validate_event'

    def test_empty_result(self):
        """Test _empty_result helper method."""
        op = StoreVectorsOperator(task_id='test')
        doc_id = UUID('550e8400-e29b-41d4-a716-446655440000')

        result = op._empty_result(doc_id)

        assert result['stored_count'] == 0
        assert result['document_id'] == '550e8400-e29b-41d4-a716-446655440000'
        assert result['status'] == 'empty'
        assert result['collection'] == ''
        assert result['document_type'] == 'unknown'

    def test_detect_document_type_excel(self):
        """Test document type detection for Excel chunks."""
        op = StoreVectorsOperator(task_id='test')

        chunks = [
            {
                'content': {'Name': 'Alice'},
                '_source': {'sheet': 'Sheet1', 'row': 1},
            }
        ]

        doc_type = op._detect_document_type(chunks)
        assert doc_type == 'excel'

    def test_detect_document_type_word(self):
        """Test document type detection for Word chunks."""
        op = StoreVectorsOperator(task_id='test')

        chunks = [
            {
                'text': 'Document content',
                'metadata': {'page_number': 1},
            }
        ]

        doc_type = op._detect_document_type(chunks)
        assert doc_type == 'word'

    def test_detect_document_type_content_sheet(self):
        """Test document type detection for 'content' sheet (semantic)."""
        op = StoreVectorsOperator(task_id='test')

        chunks = [
            {
                'content': {'text': 'Some text'},
                '_source': {'sheet': 'content'},
            }
        ]

        doc_type = op._detect_document_type(chunks)
        assert doc_type == 'word'

    def test_detect_document_type_empty(self):
        """Test document type detection for empty chunks."""
        op = StoreVectorsOperator(task_id='test')

        doc_type = op._detect_document_type([])
        assert doc_type == 'word'  # Default

    def test_generate_record_ids_with_valid_uuids(self):
        """Test record ID generation with valid UUID chunk IDs."""
        op = StoreVectorsOperator(task_id='test')

        embeddings = [
            {'chunk_id': '550e8400-e29b-41d4-a716-446655440001', 'vector': []},
            {'chunk_id': '550e8400-e29b-41d4-a716-446655440002', 'vector': []},
        ]

        ids = op._generate_record_ids(embeddings, 2)

        assert len(ids) == 2
        assert ids[0] == UUID('550e8400-e29b-41d4-a716-446655440001')
        assert ids[1] == UUID('550e8400-e29b-41d4-a716-446655440002')

    def test_generate_record_ids_with_invalid_uuids(self):
        """Test record ID generation with non-UUID chunk IDs."""
        op = StoreVectorsOperator(task_id='test')

        embeddings = [
            {'chunk_id': 'doc_t0_r1', 'vector': []},  # Not a valid UUID
        ]

        ids = op._generate_record_ids(embeddings, 1)

        assert len(ids) == 1
        # Should generate a new UUID, not raise an error
        assert isinstance(ids[0], UUID)

    def test_prepare_records_for_milvus_excel(self):
        """Test record preparation for Excel documents."""
        op = StoreVectorsOperator(task_id='test')

        chunks = [
            {
                'content': {'Name': 'Alice'},
                '_source': {'sheet': 'Sheet1', 'row': 1, 'filename': 'test.xlsx'},
            }
        ]

        records = op._prepare_records_for_milvus(chunks, 'excel')

        assert len(records) == 1
        assert isinstance(records[0], ExcelRecordProxy)
        assert records[0].content == {'Name': 'Alice'}

    def test_prepare_records_for_milvus_word(self):
        """Test record preparation for Word documents."""
        op = StoreVectorsOperator(task_id='test')

        chunks = [
            {
                'text': 'Document content',
                'metadata': {'page_number': 1},
            }
        ]

        records = op._prepare_records_for_milvus(chunks, 'word')

        assert len(records) == 1
        assert isinstance(records[0], dict)
        assert records[0]['text'] == 'Document content'

    def test_prepare_records_for_milvus_word_from_content(self):
        """Test record preparation extracts text from content dict."""
        op = StoreVectorsOperator(task_id='test')

        chunks = [
            {
                'content': {'text': 'Content from dict'},
                '_source': {'sheet': 'content'},
            }
        ]

        records = op._prepare_records_for_milvus(chunks, 'word')

        assert len(records) == 1
        assert records[0]['text'] == 'Content from dict'

    def test_store_vectors_error_handling(
        self,
        mock_context,
        sample_validate_result,
        sample_embedding_result_excel
    ):
        """Test that VectorStore errors are properly propagated."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'generate_embeddings':
                return sample_embedding_result_excel
            elif 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.store_vectors.VectorStore') as mock_store_class:
            mock_store = MagicMock()
            mock_store.create_collection = AsyncMock()
            mock_store.index_with_vectors = AsyncMock(
                side_effect=Exception('Milvus connection failed')
            )
            mock_store._disconnect = MagicMock()
            mock_store.collection_name = 'test_collection'
            mock_store_class.return_value = mock_store

            op = StoreVectorsOperator(task_id='test_store_vectors')

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

            # Verify disconnect was called even on error
            mock_store._disconnect.assert_called_once()

    def test_store_vectors_disconnects_on_success(
        self,
        mock_context,
        sample_validate_result,
        sample_embedding_result_excel
    ):
        """Test that VectorStore disconnects after successful operation."""
        def xcom_pull_side_effect(task_ids):
            if task_ids == 'generate_embeddings':
                return sample_embedding_result_excel
            elif 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.store_vectors.VectorStore') as mock_store_class:
            mock_store = MagicMock()
            mock_store.create_collection = AsyncMock()
            mock_store.index_with_vectors = AsyncMock(return_value=2)
            mock_store._disconnect = MagicMock()
            mock_store.collection_name = 'test_collection'
            mock_store_class.return_value = mock_store

            op = StoreVectorsOperator(task_id='test_store_vectors')

            result = op.execute(mock_context)

            assert result['status'] == 'success'
            # Verify disconnect was called
            mock_store._disconnect.assert_called_once()


class TestExcelRecordProxy:
    """Test suite for ExcelRecordProxy helper class."""

    def test_proxy_from_full_chunk(self):
        """Test proxy creation from complete chunk."""
        chunk = {
            'content': {'Name': 'Alice', 'Age': '25'},
            '_source': {
                'sheet': 'Sheet1',
                'row': 1,
                'filename': 'test.xlsx',
            },
            'resolved_content': {'Name': 'Alice Smith'},
        }

        proxy = ExcelRecordProxy(chunk)

        assert proxy.content == {'Name': 'Alice', 'Age': '25'}
        assert proxy.headers == ['Name', 'Age']
        assert proxy._source['sheet'] == 'Sheet1'
        assert proxy.resolved_content == {'Name': 'Alice Smith'}

    def test_proxy_from_minimal_chunk(self):
        """Test proxy creation from minimal chunk."""
        chunk = {
            'content': {'text': 'Some text'},
        }

        proxy = ExcelRecordProxy(chunk)

        assert proxy.content == {'text': 'Some text'}
        assert proxy.headers == ['text']
        # Should have default _source values
        assert proxy._source.get('sheet') == 'Unknown'
        assert proxy._source.get('row') == 0
        assert proxy._source.get('filename') == 'unknown'
        assert proxy.resolved_content is None

    def test_proxy_headers_from_non_dict_content(self):
        """Test proxy handles non-dict content."""
        chunk = {
            'content': 'Just a string',
            '_source': {'sheet': 'Sheet1'},
        }

        proxy = ExcelRecordProxy(chunk)

        assert proxy.content == 'Just a string'
        assert proxy.headers == []
