"""
Unit tests for EndeeVectorStore using mocks.
These tests verify internal logic without requiring real API credentials or internet connection.

Run with:
    pytest tests/test_vector_store.py
    # or
    python -m unittest tests/test_vector_store.py
"""

import unittest
from unittest.mock import MagicMock, patch

# Import the class to be tested
from crewai_endee import EndeeVectorStore


class TestEndeeVectorStoreUnit(unittest.TestCase):
    """
    Unit tests for EndeeVectorStore using mocks.
    No actual API keys or internet connection required.
    """

    def setUp(self):
        """Set up common mocks before each test."""
        # 1. Mock the Embedding Function
        self.mock_embedder_fn = MagicMock(return_value=[[0.1, 0.2, 0.3]])

        # 2. Mock the Sparse Encoder
        self.mock_sparse_encoder_fn = MagicMock(
            return_value=([[1, 5, 10]], [[0.5, 0.8, 0.9]])
        )

        # 3. Mock the Endee client and Index
        self.mock_endee_instance = MagicMock()
        self.mock_index_instance = MagicMock()

        # Default behavior: get_index works immediately (for save/search tests)
        self.mock_endee_instance.get_index.return_value = self.mock_index_instance
        self.mock_endee_instance.create_index.return_value = None

        # Patch the internal factory function in CrewAI to support "fake" provider
        self.patcher_factory = patch(
            "crewai.rag.embeddings.factory.build_embedder_from_dict",
            return_value=self.mock_embedder_fn,
        )

        # Patch Endee client
        self.patcher_endee = patch(
            "crewai_endee.vector_store.Endee", return_value=self.mock_endee_instance
        )

        # Patch Sparse utils
        self.patcher_sparse_get = patch(
            "crewai_endee.vector_store.get_sparse_encoder",
            return_value=self.mock_sparse_encoder_fn,
        )
        self.patcher_vocab = patch(
            "crewai_endee.vector_store.get_model_vocab_size", return_value=30522
        )

        # Start patches
        self.mock_build_factory = self.patcher_factory.start()
        self.mock_Endee_cls = self.patcher_endee.start()
        self.mock_get_sparse = self.patcher_sparse_get.start()
        self.mock_get_vocab = self.patcher_vocab.start()

    def tearDown(self):
        """Stop all patches."""
        self.patcher_factory.stop()
        self.patcher_endee.stop()
        self.patcher_sparse_get.stop()
        self.patcher_vocab.stop()

    def test_initialization_dense_only(self):
        """Test initializing in dense-only mode (default)."""
        self.mock_endee_instance.get_index.side_effect = [
            Exception("Index not found"),
            self.mock_index_instance,
        ]

        store = EndeeVectorStore(
            type="test_dense",
            embedder_config={"provider": "fake"},
            api_token="fake_token",
        )

        self.assertFalse(store._hybrid)
        store.ensure_index()

        # Verify creation was attempted
        self.mock_endee_instance.create_index.assert_called_once()
        call_kwargs = self.mock_endee_instance.create_index.call_args[1]
        self.assertNotIn("sparse_dim", call_kwargs)

    def test_initialization_hybrid_auto_detect(self):
        """Test initializing hybrid mode by providing just the model name."""
        self.mock_endee_instance.get_index.side_effect = [
            Exception("Index not found"),
            self.mock_index_instance,
        ]

        store = EndeeVectorStore(
            type="test_hybrid",
            embedder_config={"provider": "fake"},
            sparse_model_name="splade_pp",
        )

        self.assertTrue(store._hybrid)
        store.ensure_index()

        # Verify creation was attempted with sparse params
        self.mock_endee_instance.create_index.assert_called_once()
        call_kwargs = self.mock_endee_instance.create_index.call_args[1]
        self.assertIn("sparse_dim", call_kwargs)
        self.assertEqual(call_kwargs["sparse_dim"], 30522)

    def test_initialization_hybrid_explicit_dim(self):
        """Test initializing hybrid mode by providing just dimensions."""
        store = EndeeVectorStore(
            type="test_hybrid_dim",
            embedder_config={"provider": "fake"},
            sparse_dim=1000,
        )

        self.assertTrue(store._hybrid)
        self.assertEqual(store._sparse_model_name, "splade_pp")
        self.assertEqual(store._sparse_dim, 1000)

    def test_save_dense(self):
        """Test saving a document in dense mode."""
        store = EndeeVectorStore(type="test_save", embedder_config={"provider": "fake"})

        # Ensure initialized (uses default success mock from setUp)
        store.ensure_index()

        metadata = {"author": "me"}
        store.save("hello world", metadata)

        # 1. Verify embedder called
        self.mock_embedder_fn.assert_called_with(["hello world"])

        # 2. Verify upsert payload
        self.mock_index_instance.upsert.assert_called_once()
        args, _ = self.mock_index_instance.upsert.call_args
        event = args[0][0]

        self.assertEqual(event["vector"], [0.1, 0.2, 0.3])
        self.assertEqual(event["meta"]["author"], "me")
        self.assertNotIn("sparse_indices", event)

    def test_save_hybrid(self):
        """Test saving a document in hybrid mode."""
        store = EndeeVectorStore(
            type="test_save_hybrid",
            embedder_config={"provider": "fake"},
            sparse_model_name="splade_pp",
        )
        store.ensure_index()

        store.save("hello hybrid", {})

        # 1. Verify sparse encoder called
        self.mock_sparse_encoder_fn.assert_called_with(["hello hybrid"])

        # 2. Verify upsert payload
        self.mock_index_instance.upsert.assert_called_once()
        event = self.mock_index_instance.upsert.call_args[0][0][0]

        self.assertIn("sparse_indices", event)
        self.assertEqual(event["sparse_indices"], [1, 5, 10])

    def test_search_processing(self):
        """Test search query processing and result parsing."""
        store = EndeeVectorStore(
            type="test_search", embedder_config={"provider": "fake"}
        )
        store.ensure_index()

        # Mock query return from Endee
        self.mock_index_instance.query.return_value = [
            {
                "id": "123",
                "similarity": 0.95,
                "meta": {"value": "search result content", "extra": "data"},
            }
        ]

        results = store.search("find me", limit=5)

        # 1. Verify query call
        self.mock_index_instance.query.assert_called_once()
        kwargs = self.mock_index_instance.query.call_args[1]
        self.assertEqual(kwargs["vector"], [0.1, 0.2, 0.3])
        self.assertEqual(kwargs["top_k"], 5)

        # 2. Verify result formatting for CrewAI
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["content"], "search result content")
        self.assertEqual(results[0]["score"], 0.95)

    def test_reset(self):
        """Test resetting the index."""
        store = EndeeVectorStore(
            type="test_reset", embedder_config={"provider": "fake"}
        )
        store.ensure_index()

        store.reset()

        self.mock_endee_instance.delete_index.assert_called_with("test_reset")
        self.assertIsNone(store._vector_index)
        self.assertFalse(store._initialized)

    def test_text_normalization(self):
        """Test that text is truncated if it exceeds MAX_LENGTH_BYTES."""
        store = EndeeVectorStore(type="test_norm", embedder_config={"provider": "fake"})

        long_string = "a" * 10000
        normalized = store._normalize_text(long_string)

        self.assertLessEqual(len(normalized.encode("utf-8")), 8192)
        self.assertEqual(len(normalized), 8192)

    def test_search_with_filtering(self):
        """Test that metadata filters are passed correctly to Endee."""
        store = EndeeVectorStore(
            type="test_filter", embedder_config={"provider": "fake"}
        )
        store.ensure_index()

        # Mock query return
        self.mock_index_instance.query.return_value = []

        # Perform search with a filter
        my_filter = {"author": {"$eq": "me"}}
        store.search("query", filter=my_filter)

        # Verify the filter was passed to the backend
        kwargs = self.mock_index_instance.query.call_args[1]
        self.assertEqual(kwargs["filter"], my_filter)

    def test_search_score_threshold(self):
        """Test that results below the score threshold are filtered out."""
        store = EndeeVectorStore(
            type="test_threshold", embedder_config={"provider": "fake"}
        )
        store.ensure_index()

        # Mock returns one high score and one low score
        self.mock_index_instance.query.return_value = [
            {"similarity": 0.9, "meta": {"value": "good match"}},
            {"similarity": 0.1, "meta": {"value": "bad match"}},
        ]

        # Search with a threshold of 0.5
        results = store.search("query", score_threshold=0.5)

        # Should only keep the 0.9 result
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["content"], "good match")

    def test_search_error_handling(self):
        """Test that search returns empty list instead of crashing on backend error."""
        store = EndeeVectorStore(
            type="test_search_error", embedder_config={"provider": "fake"}
        )
        store.ensure_index()

        # Make the backend fail
        self.mock_index_instance.query.side_effect = Exception("DB Connection Failed")

        # Should return empty list, not raise Exception
        results = store.search("query")
        self.assertEqual(results, [])

    def test_sparse_indices_out_of_bounds(self):
        """Test that a ValueError is raised if the model generates indices larger than sparse_dim."""
        # Setup hybrid store with small dimension
        store = EndeeVectorStore(
            type="test_bounds",
            embedder_config={"provider": "fake"},
            sparse_dim=10,  # Tiny vocab size
        )
        store.ensure_index()

        # Mock encoder returning an index (100) that is >= sparse_dim (10)
        self.mock_sparse_encoder_fn.return_value = ([[100]], [[0.5]])

        # Save should fail validation
        with self.assertRaises(ValueError) as cm:
            store.save("text", {})

        self.assertIn("Sparse indices out of bounds", str(cm.exception))


if __name__ == "__main__":
    unittest.main()
