import unittest
from tests.setup_class import EndeeLangChainTestSetup
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))

from langchain_endee.vectorstores import EndeeVectorStore

# Use the new langchain-huggingface package to avoid deprecation warning
try:
    from langchain_huggingface import HuggingFaceEmbeddings
except ImportError:
    # Fallback to old import if new package not installed
    from langchain_community.embeddings import HuggingFaceEmbeddings

class TestEndeeQueryAndFilter(EndeeLangChainTestSetup):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.embed_model = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={"device": "cpu"}
        )
        
        # Use a unique index name for this test class to avoid conflicts
        cls.query_test_index = f"query_filter_{cls.test_index_name}"
        cls.test_indexes.add(cls.query_test_index)
        
        # Create vector store with force_recreate to ensure clean state
        cls.vector_store = EndeeVectorStore(
            embedding=cls.embed_model,
            api_token=cls.endee_api_token,
            index_name=cls.query_test_index,
            dimension=cls.dimension,
            space_type=cls.space_type,
            force_recreate=True  # Ensure clean slate
        )
        # Add test documents
        cls.document_ids = cls.vector_store.add_texts(
            texts=cls.test_texts, 
            metadatas=cls.test_metadatas
        )
    
    @classmethod
    def tearDownClass(cls):
        """Clean up the test index"""
        try:
            cls.nd.delete_index(name=cls.query_test_index)
            print(f"Cleaned up query test index: {cls.query_test_index}")
        except Exception as e:
            print(f"Error cleaning up query test index: {e}")
        super().tearDownClass()

    def test_basic_query(self):
        """Test basic similarity search without filters"""
        query = "What is Python?"
        results = self.vector_store.similarity_search(query, k=2)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "No results returned for basic query")
        python_found = any("python" in doc.page_content.lower() for doc in results)
        self.assertTrue(python_found, "No Python-related content found in results")

    def test_basic_query_with_score(self):
        """Test similarity search with scores"""
        query = "What is Python?"
        results = self.vector_store.similarity_search_with_score(query, k=2)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "No results returned")
        
        # Verify results are tuples of (Document, score)
        for doc, score in results:
            self.assertIsInstance(score, float)
            self.assertGreaterEqual(score, 0.0)
            self.assertLessEqual(score, 1.0)

    def test_single_filter_eq_operator(self):
        """Test filtering with $eq operator"""
        query = "What programming languages are mentioned?"
        filter_dict = [{"category": {"$eq": "programming"}}]
        results = self.vector_store.similarity_search(query, k=3, filter=filter_dict)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "No results with category=programming filter")
        
        # Check all results match the filter
        for doc in results:
            category = doc.metadata.get("category", "")
            if isinstance(category, str):
                self.assertEqual(category.lower(), "programming", 
                               f"Expected 'programming', got '{category}'")
            elif isinstance(category, list):
                self.assertIn("programming", [c.lower() for c in category],
                            f"Expected 'programming' in list, got {category}")

    def test_multiple_filters_and_operator(self):
        """Test multiple filter conditions (AND logic)"""
        query = "Tell me about AI"
        filter_dict = [{"category": {"$eq": "ai"}}, {"difficulty": {"$eq": "advanced"}}]
        results = self.vector_store.similarity_search(query, k=3, filter=filter_dict)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "No results with ai + advanced filters")
        
        for doc in results:
            category = doc.metadata.get("category", "")
            difficulty = doc.metadata.get("difficulty", "")
            
            # Check category
            if isinstance(category, str):
                self.assertEqual(category.lower(), "ai", "Category filter failed")
            elif isinstance(category, list):
                self.assertIn("ai", [c.lower() for c in category], "Category filter failed")
            
            # Check difficulty
            self.assertEqual(difficulty.lower(), "advanced", "Difficulty filter failed")

    def test_in_operator_filter(self):
        """Test filtering with $in operator"""
        query = "Tell me about programming languages"
        filter_dict = [{"difficulty": {"$in": ["intermediate", "advanced"]}}]
        results = self.vector_store.similarity_search(query, k=5, filter=filter_dict)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "No results with $in filter")
        
        for doc in results:
            difficulty = doc.metadata.get("difficulty", "").lower()
            self.assertIn(difficulty, ["intermediate", "advanced"],
                         f"Difficulty '{difficulty}' not in expected values")

    def test_in_operator_with_list_values(self):
        """Test $in operator when matching against list values in metadata"""
        query = "Tell me about programming languages"
        # This should match documents where category is a list containing "programming"
        filter_dict = [{"category": {"$in": ["programming"]}}]
        results = self.vector_store.similarity_search(query, k=5, filter=filter_dict)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "No results with $in filter on list values")
        
        for doc in results:
            category = doc.metadata.get("category", "")
            if isinstance(category, str):
                self.assertEqual(category.lower(), "programming")
            elif isinstance(category, list):
                self.assertIn("programming", [c.lower() for c in category])

    def test_range_operator_filter(self):
        """Test filtering with $range operator if supported"""
        # Note: This test assumes you have numeric metadata
        # Skipping if not applicable to your test data
        # Example: filter_dict = [{"score": {"$range": [70, 95]}}]
        self.skipTest("Range operator test skipped - no numeric metadata in test data")

    def test_invalid_filter_field(self):
        """Test behavior with non-existent filter field"""
        query = "What will I get?"
        filter_dict = [{"non_existent_field": {"$eq": "something"}}]
        results = self.vector_store.similarity_search(query, k=2, filter=filter_dict)
        self.assertIsNotNone(results)
        # Should return empty results when filtering on non-existent field
        self.assertEqual(len(results), 0, "Expected no results for invalid filter field")

    def test_no_results_filter_combination(self):
        """Test filter combination that should return no results"""
        query = "Tell me about programming languages"
        # ai + beginner combination doesn't exist in test data
        filter_dict = [{"category": {"$eq": "ai"}}, {"difficulty": {"$eq": "beginner"}}]
        results = self.vector_store.similarity_search(query, k=3, filter=filter_dict)
        self.assertIsNotNone(results)
        self.assertEqual(len(results), 0, "Expected no results for this filter combination")

    def test_filter_with_different_k_values(self):
        """Test that k parameter works correctly with filters"""
        query = "Tell me about databases"
        filter_dict = [{"category": {"$eq": "database"}}]
        
        # Request 1 result
        results_k1 = self.vector_store.similarity_search(query, k=1, filter=filter_dict)
        self.assertLessEqual(len(results_k1), 1)
        
        # Request 5 results
        results_k5 = self.vector_store.similarity_search(query, k=5, filter=filter_dict)
        # Should get at least as many results (up to available matching docs)
        self.assertGreaterEqual(len(results_k5), len(results_k1))

    def test_similarity_search_by_vector(self):
        """Test similarity search using embedding vector directly"""
        query = "What is Python?"
        # Generate embedding for the query
        query_embedding = self.embed_model.embed_query(query)
        
        # Search using the vector
        results = self.vector_store.similarity_search_by_vector(query_embedding, k=2)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)
        
        # Should get similar results as text search
        python_found = any("python" in doc.page_content.lower() for doc in results)
        self.assertTrue(python_found, "No Python-related content found in vector search results")

    def test_similarity_search_by_vector_with_score(self):
        """Test similarity search by vector with scores"""
        query = "What is Python?"
        query_embedding = self.embed_model.embed_query(query)
        
        results = self.vector_store.similarity_search_by_vector_with_score(
            query_embedding, k=2
        )
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)
        
        # Verify results are tuples of (Document, score)
        for doc, score in results:
            self.assertIsInstance(score, float)
            self.assertGreaterEqual(score, 0.0)
            self.assertLessEqual(score, 1.0)

    def test_similarity_search_by_vector_with_filter(self):
        """Test similarity search by vector with filter"""
        query = "Tell me about programming"
        query_embedding = self.embed_model.embed_query(query)
        filter_dict = [{"category": {"$eq": "programming"}}]
        
        results = self.vector_store.similarity_search_by_vector(
            query_embedding, k=3, filter=filter_dict
        )
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)
        
        # Verify all results match filter
        for doc in results:
            category = doc.metadata.get("category", "")
            if isinstance(category, str):
                self.assertEqual(category.lower(), "programming")
            elif isinstance(category, list):
                self.assertIn("programming", [c.lower() for c in category])

    def test_custom_ef_parameter(self):
        """Test custom ef (search quality) parameter"""
        query = "What is Python?"

        # Try with different ef values
        results_ef64 = self.vector_store.similarity_search(query, k=2, ef=64)
        results_ef256 = self.vector_store.similarity_search(query, k=2, ef=256)

        # Both should return results
        self.assertGreater(len(results_ef64), 0)
        self.assertGreater(len(results_ef256), 0)

        # Higher ef might return better quality results (but not guaranteed to be different)
        # Just verify both work without errors

    def test_prefilter_cardinality_threshold(self):
        """Test prefilter_cardinality_threshold parameter on all four search methods"""
        query = "What is Python?"
        query_embedding = self.embed_model.embed_query(query)
        filter_dict = [{"category": {"$eq": "programming"}}]

        # similarity_search
        results = self.vector_store.similarity_search(
            query, k=2, filter=filter_dict, prefilter_cardinality_threshold=5_000
        )
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

        # similarity_search_with_score
        results_score = self.vector_store.similarity_search_with_score(
            query, k=2, filter=filter_dict, prefilter_cardinality_threshold=5_000
        )
        self.assertIsNotNone(results_score)
        self.assertGreater(len(results_score), 0)
        for doc, score in results_score:
            self.assertIsInstance(score, float)

        # similarity_search_by_vector
        results_vec = self.vector_store.similarity_search_by_vector(
            query_embedding, k=2, filter=filter_dict, prefilter_cardinality_threshold=5_000
        )
        self.assertIsNotNone(results_vec)
        self.assertGreater(len(results_vec), 0)

        # similarity_search_by_vector_with_score
        results_vec_score = self.vector_store.similarity_search_by_vector_with_score(
            query_embedding, k=2, filter=filter_dict, prefilter_cardinality_threshold=5_000
        )
        self.assertIsNotNone(results_vec_score)
        self.assertGreater(len(results_vec_score), 0)

    def test_filter_boost_percentage(self):
        """Test filter_boost_percentage parameter on all four search methods"""
        query = "What is Python?"
        query_embedding = self.embed_model.embed_query(query)
        filter_dict = [{"category": {"$eq": "programming"}}]

        # similarity_search
        results = self.vector_store.similarity_search(
            query, k=2, filter=filter_dict, filter_boost_percentage=25
        )
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

        # similarity_search_with_score
        results_score = self.vector_store.similarity_search_with_score(
            query, k=2, filter=filter_dict, filter_boost_percentage=25
        )
        self.assertIsNotNone(results_score)
        self.assertGreater(len(results_score), 0)
        for doc, score in results_score:
            self.assertIsInstance(score, float)

        # similarity_search_by_vector
        results_vec = self.vector_store.similarity_search_by_vector(
            query_embedding, k=2, filter=filter_dict, filter_boost_percentage=25
        )
        self.assertIsNotNone(results_vec)
        self.assertGreater(len(results_vec), 0)

        # similarity_search_by_vector_with_score
        results_vec_score = self.vector_store.similarity_search_by_vector_with_score(
            query_embedding, k=2, filter=filter_dict, filter_boost_percentage=25
        )
        self.assertIsNotNone(results_vec_score)
        self.assertGreater(len(results_vec_score), 0)

    def test_prefilter_and_boost_combined(self):
        """Test prefilter_cardinality_threshold and filter_boost_percentage together"""
        query = "Tell me about databases"
        filter_dict = [{"category": {"$eq": "database"}}]

        results = self.vector_store.similarity_search(
            query,
            k=3,
            filter=filter_dict,
            prefilter_cardinality_threshold=5_000,
            filter_boost_percentage=30,
        )
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

        # All results should match the filter
        for doc in results:
            category = doc.metadata.get("category", "")
            if isinstance(category, str):
                self.assertEqual(category.lower(), "database")
            elif isinstance(category, list):
                self.assertIn("database", [c.lower() for c in category])

    def test_prefilter_boundary_values(self):
        """Test prefilter_cardinality_threshold and filter_boost_percentage at boundary values"""
        query = "What is Python?"
        filter_dict = [{"category": {"$eq": "programming"}}]

        # Min boundary for prefilter_cardinality_threshold (1_000)
        results_min = self.vector_store.similarity_search(
            query, k=2, filter=filter_dict, prefilter_cardinality_threshold=1_000
        )
        self.assertIsNotNone(results_min)

        # Large value for prefilter_cardinality_threshold (1_000_000)
        results_max = self.vector_store.similarity_search(
            query, k=2, filter=filter_dict, prefilter_cardinality_threshold=1_000_000
        )
        self.assertIsNotNone(results_max)

        # filter_boost_percentage = 0 (no boost)
        results_no_boost = self.vector_store.similarity_search(
            query, k=2, filter=filter_dict, filter_boost_percentage=0
        )
        self.assertIsNotNone(results_no_boost)

        # filter_boost_percentage = 100 (double candidate pool)
        results_full_boost = self.vector_store.similarity_search(
            query, k=2, filter=filter_dict, filter_boost_percentage=100
        )
        self.assertIsNotNone(results_full_boost)

    def test_filter_params_without_filter(self):
        """Test that prefilter_cardinality_threshold and filter_boost_percentage
        work gracefully when no filter is provided (parameters should be ignored)"""
        query = "What is Python?"

        results = self.vector_store.similarity_search(
            query,
            k=2,
            prefilter_cardinality_threshold=5_000,
            filter_boost_percentage=20,
        )
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

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