"""Unit tests for vector_store module."""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch, call
from uuid import uuid4, UUID

from pymilvus import Collection, CollectionSchema

from src.core.exceptions import IndexingError
from src.extraction.models import ExtractedRecord
from src.knowledge.vector_store import MilvusEntity, VectorStore


class TestMilvusEntity:
    """Tests for MilvusEntity dataclass."""

    def test_milvus_entity_creation(self):
        """Test creating MilvusEntity with all fields."""
        vector = [0.1] * 3072
        record_id = str(uuid4())
        document_id = str(uuid4())

        entity = MilvusEntity(
            vector=vector,
            record_id=record_id,
            document_id=document_id,
            filename="test.xlsx",
            sheet_name="Sheet1",
            row_number=10,
            text_content="Item Code: ABC. From: test.xlsx, Sheet: Sheet1, Row: 10",
            has_symbols=False,
        )

        assert len(entity.vector) == 3072
        assert entity.record_id == record_id
        assert entity.document_id == document_id
        assert entity.filename == "test.xlsx"
        assert entity.sheet_name == "Sheet1"
        assert entity.row_number == 10
        assert "Item Code: ABC" in entity.text_content
        assert entity.has_symbols is False


class TestVectorStore:
    """Tests for VectorStore class."""

    @pytest.fixture
    def mock_embedding_service(self):
        """Create mock EmbeddingService."""
        with patch("src.knowledge.vector_store.EmbeddingService") as mock:
            service = mock.return_value
            service.embed_batch = AsyncMock(return_value=[[0.1] * 3072, [0.2] * 3072])
            yield service

    @pytest.fixture
    def mock_milvus_connections(self):
        """Mock pymilvus connections."""
        with patch("src.knowledge.vector_store.connections") as mock:
            yield mock

    @pytest.fixture
    def mock_milvus_collection(self):
        """Mock pymilvus Collection."""
        with patch("src.knowledge.vector_store.Collection") as mock:
            yield mock

    @pytest.fixture
    def mock_milvus_utility(self):
        """Mock pymilvus utility."""
        with patch("src.knowledge.vector_store.utility") as mock:
            yield mock

    @pytest.fixture
    def vector_store(self, mock_embedding_service):
        """Create VectorStore with mocked dependencies."""
        return VectorStore()

    def test_vector_store_initialization(self, vector_store):
        """Test VectorStore initialization."""
        assert vector_store.collection_name == "dsol_records"
        assert vector_store.vector_dim == 3072
        assert not vector_store._connected

    def test_connect(self, vector_store, mock_milvus_connections):
        """Test Milvus connection."""
        vector_store._connect()

        assert vector_store._connected
        mock_milvus_connections.connect.assert_called_once_with(
            alias="default",
            host="localhost",
            port=19530,
        )

    def test_connect_idempotent(self, vector_store, mock_milvus_connections):
        """Test that connect is idempotent."""
        vector_store._connect()
        vector_store._connect()

        # Should only connect once
        assert mock_milvus_connections.connect.call_count == 1

    def test_disconnect(self, vector_store, mock_milvus_connections):
        """Test Milvus disconnection."""
        vector_store._connected = True
        vector_store._disconnect()

        assert not vector_store._connected
        mock_milvus_connections.disconnect.assert_called_once_with("default")

    @pytest.mark.asyncio
    async def test_create_collection_new(
        self,
        vector_store,
        mock_milvus_connections,
        mock_milvus_collection,
        mock_milvus_utility,
    ):
        """Test creating new collection."""
        mock_milvus_utility.has_collection.return_value = False

        mock_collection_instance = MagicMock()
        mock_milvus_collection.return_value = mock_collection_instance

        await vector_store.create_collection()

        # Verify collection creation
        mock_milvus_utility.has_collection.assert_called_once_with("dsol_records")
        mock_milvus_collection.assert_called_once()

        # Verify index creation
        mock_collection_instance.create_index.assert_called_once()
        index_call = mock_collection_instance.create_index.call_args
        assert index_call.kwargs["field_name"] == "vector"
        assert index_call.kwargs["index_params"]["metric_type"] == "COSINE"
        assert index_call.kwargs["index_params"]["index_type"] == "HNSW"

        # Verify collection loaded
        mock_collection_instance.load.assert_called_once()

    @pytest.mark.asyncio
    async def test_create_collection_already_exists(
        self,
        vector_store,
        mock_milvus_connections,
        mock_milvus_collection,
        mock_milvus_utility,
    ):
        """Test creating collection when it already exists."""
        mock_milvus_utility.has_collection.return_value = True

        await vector_store.create_collection()

        # Should not attempt to create
        mock_milvus_collection.assert_not_called()

    @pytest.mark.asyncio
    async def test_create_collection_failure(
        self,
        vector_store,
        mock_milvus_connections,
        mock_milvus_utility,
    ):
        """Test collection creation failure."""
        mock_milvus_utility.has_collection.side_effect = Exception("Connection failed")

        with pytest.raises(IndexingError, match="Failed to create Milvus collection"):
            await vector_store.create_collection()

    @pytest.mark.asyncio
    async def test_index_records_success(
        self,
        vector_store,
        mock_embedding_service,
        mock_milvus_connections,
        mock_milvus_collection,
    ):
        """Test successful record indexing."""
        # Create test records
        document_id = uuid4()
        records = [
            ExtractedRecord(
                content={"Item": "A"},
                headers=["Item"],
                _source={
                    "document_id": str(document_id),
                    "filename": "test.xlsx",
                    "sheet": "Sheet1",
                    "table_id": 1,
                    "row": 1,
                    "col_range": "A:A",
                },
            ),
            ExtractedRecord(
                content={"Item": "B"},
                headers=["Item"],
                _source={
                    "document_id": str(document_id),
                    "filename": "test.xlsx",
                    "sheet": "Sheet1",
                    "table_id": 1,
                    "row": 2,
                    "col_range": "A:A",
                },
            ),
        ]

        record_ids = [uuid4(), uuid4()]

        # Mock collection
        mock_collection_instance = MagicMock()
        mock_milvus_collection.return_value = mock_collection_instance

        # Index records
        count = await vector_store.index_records(records, document_id, record_ids)

        # Verify results
        assert count == 2
        mock_embedding_service.embed_batch.assert_called_once()
        mock_collection_instance.upsert.assert_called_once()

    @pytest.mark.asyncio
    async def test_index_records_mismatched_lengths(self, vector_store):
        """Test error when records and record_ids lengths don't match."""
        records = [MagicMock()]
        record_ids = [uuid4(), uuid4()]

        with pytest.raises(ValueError, match="Records count.*must match"):
            await vector_store.index_records(records, uuid4(), record_ids)

    @pytest.mark.asyncio
    async def test_index_records_empty_list(self, vector_store):
        """Test handling of empty record list."""
        count = await vector_store.index_records([], uuid4(), [])

        assert count == 0

    @pytest.mark.asyncio
    async def test_index_records_batch_processing(
        self,
        vector_store,
        mock_embedding_service,
        mock_milvus_connections,
        mock_milvus_collection,
    ):
        """Test that records are processed in batches of 100."""
        # Create 150 records to test batching
        document_id = uuid4()
        records = []
        record_ids = []

        for i in range(150):
            records.append(
                ExtractedRecord(
                    content={"Item": f"Item{i}"},
                    headers=["Item"],
                    _source={
                        "document_id": str(document_id),
                        "filename": "test.xlsx",
                        "sheet": "Sheet1",
                        "table_id": 1,
                        "row": i + 1,
                        "col_range": "A:A",
                    },
                )
            )
            record_ids.append(uuid4())

        # Mock embedding service to return appropriate vectors
        mock_embedding_service.embed_batch.return_value = [[0.1] * 3072] * 100

        # Mock collection
        mock_collection_instance = MagicMock()
        mock_milvus_collection.return_value = mock_collection_instance

        # Index records
        count = await vector_store.index_records(records, document_id, record_ids)

        # Verify batching: 2 calls (100 + 50)
        assert count == 150
        assert mock_embedding_service.embed_batch.call_count == 2

        # Verify first batch size
        first_call = mock_embedding_service.embed_batch.call_args_list[0]
        assert len(first_call[0][0]) == 100

        # Verify second batch size
        second_call = mock_embedding_service.embed_batch.call_args_list[1]
        assert len(second_call[0][0]) == 50

    def test_prepare_milvus_entities(self, vector_store):
        """Test entity preparation from records."""
        document_id = uuid4()
        records = [
            ExtractedRecord(
                content={"Item": "A"},
                headers=["Item"],
                _source={
                    "document_id": str(document_id),
                    "filename": "test.xlsx",
                    "sheet": "Sheet1",
                    "table_id": 1,
                    "row": 1,
                    "col_range": "A:A",
                },
            ),
        ]

        embeddings = [[0.1] * 3072]
        record_ids = [uuid4()]

        entities = vector_store._prepare_milvus_entities(
            records, embeddings, record_ids, document_id
        )

        assert len(entities) == 1
        entity = entities[0]
        assert len(entity.vector) == 3072
        assert entity.record_id == str(record_ids[0])
        assert entity.document_id == str(document_id)
        assert entity.filename == "test.xlsx"
        assert entity.sheet_name == "Sheet1"
        assert entity.row_number == 1
        assert "Item: A" in entity.text_content
        assert entity.has_symbols is False

    def test_prepare_milvus_entities_with_symbols(self, vector_store):
        """Test entity preparation for record with resolved content."""
        document_id = uuid4()
        record = ExtractedRecord(
            content={"Screen": "○"},
            headers=["Screen"],
            _source={
                "document_id": str(document_id),
                "filename": "test.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 1,
                "col_range": "A:A",
            },
        )

        # Add resolved_content
        record.resolved_content = {
            "Screen": "○",
            "Screen_resolved": "Applicable/Yes",
        }

        embeddings = [[0.1] * 3072]
        record_ids = [uuid4()]

        entities = vector_store._prepare_milvus_entities(
            [record], embeddings, record_ids, document_id
        )

        assert len(entities) == 1
        assert entities[0].has_symbols is True
        assert "Screen_resolved: Applicable/Yes" in entities[0].text_content

    def test_prepare_milvus_entities_text_truncation(self, vector_store):
        """Test that long text content is truncated to 2000 chars."""
        document_id = uuid4()
        long_value = "X" * 2000  # Create long value

        record = ExtractedRecord(
            content={"LongField": long_value},
            headers=["LongField"],
            _source={
                "document_id": str(document_id),
                "filename": "test.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 1,
                "col_range": "A:A",
            },
        )

        embeddings = [[0.1] * 3072]
        record_ids = [uuid4()]

        entities = vector_store._prepare_milvus_entities(
            [record], embeddings, record_ids, document_id
        )

        # Text should be truncated to 2000 chars
        assert len(entities[0].text_content) == 2000
        assert entities[0].text_content.endswith("...")

    @pytest.mark.asyncio
    async def test_upsert_vectors_success(self, vector_store, mock_milvus_collection):
        """Test successful vector upsert."""
        entities = [
            MilvusEntity(
                vector=[0.1] * 3072,
                record_id=str(uuid4()),
                document_id=str(uuid4()),
                filename="test.xlsx",
                sheet_name="Sheet1",
                row_number=1,
                text_content="Test content",
                has_symbols=False,
            ),
        ]

        mock_collection_instance = MagicMock()
        mock_milvus_collection.return_value = mock_collection_instance

        await vector_store._upsert_vectors(entities)

        # Verify upsert called
        mock_collection_instance.upsert.assert_called_once()
        call_data = mock_collection_instance.upsert.call_args[0][0]

        # Verify data structure (columnar format)
        assert len(call_data) == 8  # 8 fields
        assert len(call_data[0]) == 1  # 1 entity
        assert call_data[0][0] == [0.1] * 3072  # vector

    @pytest.mark.asyncio
    async def test_upsert_vectors_retry_logic(
        self, vector_store, mock_milvus_collection
    ):
        """Test retry logic on transient errors."""
        entities = [
            MilvusEntity(
                vector=[0.1] * 3072,
                record_id=str(uuid4()),
                document_id=str(uuid4()),
                filename="test.xlsx",
                sheet_name="Sheet1",
                row_number=1,
                text_content="Test content",
                has_symbols=False,
            ),
        ]

        mock_collection_instance = MagicMock()
        mock_collection_instance.upsert.side_effect = [
            Exception("Timeout"),
            Exception("Network error"),
            None,  # Success on 3rd try
        ]
        mock_milvus_collection.return_value = mock_collection_instance

        # Should succeed after retries
        await vector_store._upsert_vectors(entities, max_retries=3)

        assert mock_collection_instance.upsert.call_count == 3

    @pytest.mark.asyncio
    async def test_upsert_vectors_permanent_failure(
        self, vector_store, mock_milvus_collection
    ):
        """Test IndexingError raised after max retries."""
        entities = [
            MilvusEntity(
                vector=[0.1] * 3072,
                record_id=str(uuid4()),
                document_id=str(uuid4()),
                filename="test.xlsx",
                sheet_name="Sheet1",
                row_number=1,
                text_content="Test content",
                has_symbols=False,
            ),
        ]

        mock_collection_instance = MagicMock()
        mock_collection_instance.upsert.side_effect = Exception("Permanent error")
        mock_milvus_collection.return_value = mock_collection_instance

        with pytest.raises(IndexingError, match="Failed to upsert vectors"):
            await vector_store._upsert_vectors(entities, max_retries=3)

        assert mock_collection_instance.upsert.call_count == 3

    def test_context_manager(self, vector_store, mock_milvus_connections):
        """Test VectorStore as context manager."""
        with vector_store:
            assert vector_store._connected
            mock_milvus_connections.connect.assert_called_once()

        assert not vector_store._connected
        mock_milvus_connections.disconnect.assert_called_once()


class TestRecordJsonMetadata:
    """Tests for record_json metadata field (LARGE TABLE pipeline)."""

    @pytest.fixture
    def vector_store(self):
        """Create VectorStore instance."""
        with patch("src.knowledge.vector_store.EmbeddingService"):
            return VectorStore()

    def test_record_json_added_for_standard_excel_records(self, vector_store):
        """Test that record_json is added for standard Excel records (LARGE TABLE pipeline)."""
        document_id = uuid4()
        record = ExtractedRecord(
            content={"Name": "Alice", "Age": "30", "Department": "Engineering"},
            headers=["Name", "Age", "Department"],
            _source={
                "document_id": str(document_id),
                "filename": "employees.xlsx",
                "sheet": "Staff",
                "table_id": 1,
                "row": 5,
                "col_range": "A:C",
            },
        )

        entities = vector_store._prepare_milvus_entities(
            [record],
            [[0.1] * 3072],
            str(document_id),
            "employees.xlsx",
            "excel",
        )

        assert len(entities) == 1
        entity = entities[0]

        # Verify record_json is present
        assert "record_json" in entity.metadata
        assert entity.metadata["record_json"] == record.content
        assert entity.metadata["record_json"]["Name"] == "Alice"
        assert entity.metadata["record_json"]["Age"] == "30"
        assert entity.metadata["record_json"]["Department"] == "Engineering"

        # Verify other metadata still present
        assert entity.metadata["sheet_name"] == "Staff"
        assert entity.metadata["row_number"] == 5

    def test_record_json_not_added_for_detailed_pipeline(self, vector_store):
        """Test that record_json is NOT added for DETAILED pipeline records (images, flowcharts, shapes)."""
        document_id = uuid4()
        record = ExtractedRecord(
            content={"type": "image", "description": "Company logo"},
            headers=["type", "description"],
            _source={
                "document_id": str(document_id),
                "filename": "doc.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 1,
                "col_range": "A:B",
            },
        )

        # Add _legacy_metadata to simulate DETAILED pipeline record
        record._legacy_metadata = {
            "element_type": "Image",
            "full_metadata": {
                "image_number": 1,
                "position_key": "img_001",
            },
        }

        entities = vector_store._prepare_milvus_entities(
            [record],
            [[0.1] * 3072],
            str(document_id),
            "doc.xlsx",
            "excel",
        )

        assert len(entities) == 1
        entity = entities[0]

        # Verify record_json is NOT present
        assert "record_json" not in entity.metadata

        # Verify DETAILED pipeline metadata is present
        assert entity.metadata["element_type"] == "Image"
        assert "row_number" not in entity.metadata  # Removed for DETAILED records

    def test_record_json_truncation_for_large_records(self, vector_store):
        """Test that large record_json is truncated to prevent Milvus errors."""
        document_id = uuid4()

        # Create a large record (>50KB when serialized)
        large_content = {f"Field_{i}": "X" * 1000 for i in range(100)}  # ~100KB

        record = ExtractedRecord(
            content=large_content,
            headers=list(large_content.keys()),
            _source={
                "document_id": str(document_id),
                "filename": "large.xlsx",
                "sheet": "Data",
                "table_id": 1,
                "row": 10,
                "col_range": "A:CZ",
            },
        )

        entities = vector_store._prepare_milvus_entities(
            [record],
            [[0.1] * 3072],
            str(document_id),
            "large.xlsx",
            "excel",
        )

        assert len(entities) == 1
        entity = entities[0]

        # Verify record_json is truncated
        assert "record_json" in entity.metadata
        assert entity.metadata["record_json_truncated"] is True
        # Should have only first 20 fields
        assert len(entity.metadata["record_json"]) == 20

    def test_record_json_serialization_error_handling(self, vector_store):
        """Test handling of non-serializable content in record_json."""
        document_id = uuid4()

        # Create record with non-serializable content
        # (This is a test case - in practice, record.content should be serializable)
        class NonSerializable:
            pass

        record = ExtractedRecord(
            content={"Field": "Value"},  # Will be normal
            headers=["Field"],
            _source={
                "document_id": str(document_id),
                "filename": "test.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 1,
                "col_range": "A:A",
            },
        )

        # Simulate serialization error by mocking json.dumps
        with patch("src.knowledge.vector_store.json.dumps", side_effect=TypeError("Not serializable")):
            entities = vector_store._prepare_milvus_entities(
                [record],
                [[0.1] * 3072],
                str(document_id),
                "test.xlsx",
                "excel",
            )

            assert len(entities) == 1
            entity = entities[0]

            # Should store as string fallback
            assert "record_json" in entity.metadata
            assert isinstance(entity.metadata["record_json"], str)

    def test_word_documents_do_not_have_record_json(self, vector_store):
        """Test that Word documents do NOT have record_json field."""
        document_id = uuid4()

        # Word record format (dict, not ExtractedRecord)
        word_record = {
            "text": "This is a paragraph from a Word document.",
            "metadata": {
                "page_number": 1,
                "contains_table": False,
            },
        }

        entities = vector_store._prepare_milvus_entities(
            [word_record],
            [[0.1] * 3072],
            str(document_id),
            "document.docx",
            "word",
        )

        assert len(entities) == 1
        entity = entities[0]

        # Verify NO record_json for Word documents
        assert "record_json" not in entity.metadata

        # Verify Word-specific metadata is present
        assert entity.metadata["page_number"] == 1
        assert entity.metadata["contains_table"] is False

    def test_existing_metadata_fields_preserved_with_record_json(self, vector_store):
        """Test that existing metadata fields are preserved when record_json is added."""
        document_id = uuid4()
        record = ExtractedRecord(
            content={"Item": "Widget", "Status": "Active"},
            headers=["Item", "Status"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Items",
                "table_id": 1,
                "row": 42,
                "col_range": "A:B",
                "cell_colors": {"A42": "#FF0000"},  # Cell color metadata
            },
        )

        # Add resolved_content to test has_symbols
        record.resolved_content = {
            "Item": "Widget",
            "Status": "Active",
            "Status_resolved": "Currently Active",
        }

        entities = vector_store._prepare_milvus_entities(
            [record],
            [[0.1] * 3072],
            str(document_id),
            "inventory.xlsx",
            "excel",
        )

        assert len(entities) == 1
        entity = entities[0]

        # Verify record_json is present
        assert "record_json" in entity.metadata
        assert entity.metadata["record_json"]["Item"] == "Widget"

        # Verify existing metadata fields are still present
        assert entity.metadata["sheet_name"] == "Items"
        assert entity.metadata["row_number"] == 42
        assert entity.metadata["has_symbols"] is True
        assert entity.metadata["cell_colors"] == {"A42": "#FF0000"}