"""Integration tests for DomainAwareRetriever.

Note: These tests require a running Milvus instance and are skipped
if Milvus is not available.
"""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from src.knowledge.domain_aware_retriever import DomainAwareRetriever
from src.knowledge.query_term_matcher import QueryTermMatcher


@pytest.fixture
def mock_query_matcher():
    """Create a mock QueryTermMatcher."""
    matcher = AsyncMock(spec=QueryTermMatcher)
    matcher.extract_query_terms = AsyncMock(return_value={
        "matched_terms": ["revenue", "ebitda"],
        "domain_category": "finance",
        "matching_method": "keyword"
    })
    return matcher


@pytest.fixture
def mock_embedding_service():
    """Create a mock EmbeddingService."""
    service = AsyncMock()
    # Return a dummy 3072-dim vector
    service.embed = AsyncMock(return_value=[0.1] * 3072)
    return service


@pytest.fixture
def retriever(mock_query_matcher, mock_embedding_service):
    """Create DomainAwareRetriever with mocked dependencies."""
    return DomainAwareRetriever(
        query_matcher=mock_query_matcher,
        embedding_service=mock_embedding_service,
    )


class TestDomainAwareRetriever:
    """Unit tests for DomainAwareRetriever."""

    def test_build_filter_with_terms_and_category(self, retriever):
        """Should build filter with category and terms."""
        matched_terms = ["revenue", "ebitda", "profit"]
        category = "finance"

        filter_expr = retriever._build_filter(
            matched_terms, category, "tenant1", None
        )

        assert filter_expr is not None
        assert "domain_category ==" in filter_expr
        assert "finance" in filter_expr
        assert "array_contains_any" in filter_expr
        assert "revenue" in filter_expr

    def test_build_filter_with_document_id(self, retriever):
        """Should include document_id in filter when provided."""
        doc_id = "123e4567-e89b-12d3-a456-426614174000"

        filter_expr = retriever._build_filter(
            ["revenue"], "finance", "tenant1", doc_id
        )

        assert filter_expr is not None
        assert f'document_id == "{doc_id}"' in filter_expr

    def test_build_filter_general_category_skipped(self, retriever):
        """Should not filter by category when category is 'general'."""
        filter_expr = retriever._build_filter(
            ["term1"], "general", "tenant1", None
        )

        # Should only have terms filter, not category filter
        if filter_expr:
            assert "domain_category" not in filter_expr
            assert "array_contains_any" in filter_expr

    def test_build_filter_no_terms_or_category(self, retriever):
        """Should return None when no terms or category."""
        filter_expr = retriever._build_filter(
            [], "general", "tenant1", None
        )

        assert filter_expr is None

    def test_build_filter_escapes_quotes(self, retriever):
        """Should escape quotes in terms."""
        matched_terms = ["term's", 'term"s']

        filter_expr = retriever._build_filter(
            matched_terms, "finance", "tenant1", None
        )

        # Should have escaped the quotes
        assert filter_expr is not None
        # The single quote should be escaped
        assert "term" in filter_expr

    def test_build_filter_only_document_id(self, retriever):
        """Should build filter with only document_id."""
        doc_id = "123e4567-e89b-12d3-a456-426614174000"

        filter_expr = retriever._build_filter(
            [], "general", "tenant1", doc_id
        )

        assert filter_expr == f'document_id == "{doc_id}"'

    @pytest.mark.asyncio
    @patch("src.knowledge.domain_aware_retriever.connections")
    @patch("src.knowledge.domain_aware_retriever.Collection")
    async def test_retrieve_uses_filter(
        self, mock_collection_class, mock_connections,
        mock_query_matcher, mock_embedding_service
    ):
        """Should pass filter expression to Milvus search."""
        # Mock Milvus collection
        mock_collection = MagicMock()
        mock_search_results = [[]]  # Empty results
        mock_collection.search.return_value = mock_search_results
        mock_collection_class.return_value = mock_collection

        retriever = DomainAwareRetriever(
            query_matcher=mock_query_matcher,
            embedding_service=mock_embedding_service,
        )

        await retriever.retrieve("What is Q4 EBITDA?")

        # Verify search was called with a filter
        assert mock_collection.search.called
        call_kwargs = mock_collection.search.call_args.kwargs
        assert "expr" in call_kwargs
        # Should have a filter (not None) since we have matched terms
        assert call_kwargs["expr"] is not None

    @pytest.mark.asyncio
    @patch("src.knowledge.domain_aware_retriever.connections")
    @patch("src.knowledge.domain_aware_retriever.Collection")
    async def test_retrieve_formats_results(
        self, mock_collection_class, mock_connections,
        mock_query_matcher, mock_embedding_service
    ):
        """Should format Milvus results correctly."""
        # Mock Milvus search results
        mock_hit = MagicMock()
        mock_hit.entity.get.side_effect = lambda key: {
            "record_id": "rec_123",
            "document_id": "doc_456",
            "filename": "test.xlsx",
            "document_type": "excel",
            "text_content": "Revenue was $1M",
            "metadata": {"sheet_name": "Sheet1"},
            "domain_category": "finance",
            "domain_terms": ["revenue"],
        }.get(key)
        mock_hit.score = 0.95

        mock_collection = MagicMock()
        mock_collection.search.return_value = [[mock_hit]]
        mock_collection_class.return_value = mock_collection

        retriever = DomainAwareRetriever(
            query_matcher=mock_query_matcher,
            embedding_service=mock_embedding_service,
        )

        results = await retriever.retrieve("What is revenue?")

        assert len(results) == 1
        assert results[0]["record_id"] == "rec_123"
        assert results[0]["document_id"] == "doc_456"
        assert results[0]["text_content"] == "Revenue was $1M"
        assert results[0]["domain_category"] == "finance"
        assert results[0]["score"] == 0.95

    @pytest.mark.asyncio
    async def test_retrieve_with_no_matches(
        self, mock_query_matcher, mock_embedding_service
    ):
        """Should handle case when query matcher returns no terms."""
        # Configure matcher to return no terms
        mock_query_matcher.extract_query_terms = AsyncMock(return_value={
            "matched_terms": [],
            "domain_category": "general",
            "matching_method": "none"
        })

        with patch("src.knowledge.domain_aware_retriever.connections"):
            with patch("src.knowledge.domain_aware_retriever.Collection") as mock_coll:
                mock_collection = MagicMock()
                mock_collection.search.return_value = [[]]
                mock_coll.return_value = mock_collection

                retriever = DomainAwareRetriever(
                    query_matcher=mock_query_matcher,
                    embedding_service=mock_embedding_service,
                )

                results = await retriever.retrieve("generic query")

                # Should still perform search, but without filter
                assert mock_collection.search.called
                call_kwargs = mock_collection.search.call_args.kwargs
                assert call_kwargs["expr"] is None  # No filter

    @pytest.mark.asyncio
    @patch("src.knowledge.domain_aware_retriever.connections")
    @patch("src.knowledge.domain_aware_retriever.Collection")
    async def test_retrieve_respects_top_k(
        self, mock_collection_class, mock_connections,
        mock_query_matcher, mock_embedding_service
    ):
        """Should pass top_k parameter to Milvus search."""
        mock_collection = MagicMock()
        mock_collection.search.return_value = [[]]
        mock_collection_class.return_value = mock_collection

        retriever = DomainAwareRetriever(
            query_matcher=mock_query_matcher,
            embedding_service=mock_embedding_service,
        )

        await retriever.retrieve("test query", top_k=25)

        # Verify limit parameter
        call_kwargs = mock_collection.search.call_args.kwargs
        assert call_kwargs["limit"] == 25


class TestIntegration:
    """Integration tests (require running Milvus instance)."""

    @pytest.mark.integration
    @pytest.mark.asyncio
    async def test_full_retrieval_pipeline(self):
        """Test complete retrieval flow with real Milvus.

        Note: This test is marked as 'integration' and requires:
        - Running Milvus instance
        - Populated collection with domain-aware data
        """
        # This test would require actual Milvus setup
        # Skip for now - would be run in integration test suite
        pytest.skip("Requires running Milvus instance with test data")

    @pytest.mark.integration
    @pytest.mark.asyncio
    async def test_performance_with_filtering_vs_full_search(self):
        """Benchmark filtered vs full search performance.

        Note: This test is marked as 'integration' and requires:
        - Running Milvus instance
        - Large dataset (100k+ chunks)
        """
        pytest.skip("Requires running Milvus instance with large dataset")
