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

from tests.setup_class import EndeeLangChainTestSetup
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 TestRetrieverDeleteReconnect(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.retriever_test_index = f"retriever_delete_{cls.test_index_name}"
        cls.test_indexes.add(cls.retriever_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.retriever_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.retriever_test_index)
            print(f"Cleaned up retriever test index: {cls.retriever_test_index}")
        except Exception as e:
            print(f"Error cleaning up retriever test index: {e}")
        super().tearDownClass()

    def test_retriever_basic(self):
        """Test basic retriever functionality"""
        retriever = self.vector_store.as_retriever(search_kwargs={"k": 2})
        retrieved_docs = retriever.invoke("What is machine learning?")
        
        self.assertIsNotNone(retrieved_docs)
        self.assertGreater(len(retrieved_docs), 0, "Retriever returned no documents")
        self.assertLessEqual(len(retrieved_docs), 2, "Retriever returned more than k documents")
        
        found = any("machine learning" in doc.page_content.lower() for doc in retrieved_docs)
        self.assertTrue(found, "No machine learning content found in retriever results")

    def test_retriever_with_filters(self):
        """Test retriever with filter arguments"""
        retriever = self.vector_store.as_retriever(
            search_kwargs={
                "k": 3,
                "filter": [{"category": {"$eq": "programming"}}]
            }
        )
        retrieved_docs = retriever.invoke("Tell me about programming languages")
        
        self.assertIsNotNone(retrieved_docs)
        self.assertGreater(len(retrieved_docs), 0, "No documents retrieved with filter")
        
        # Verify all results match the filter
        for doc in retrieved_docs:
            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_retriever_different_search_types(self):
        """Test retriever with different search types (if supported)"""
        # Test similarity search (default)
        retriever_similarity = self.vector_store.as_retriever(
            search_type="similarity",
            search_kwargs={"k": 2}
        )
        results_similarity = retriever_similarity.invoke("Python programming")
        self.assertGreater(len(results_similarity), 0)
        
        # Test similarity_score_threshold if supported
        try:
            retriever_threshold = self.vector_store.as_retriever(
                search_type="similarity_score_threshold",
                search_kwargs={"score_threshold": 0.5, "k": 5}
            )
            results_threshold = retriever_threshold.invoke("Python programming")
            # Should return only results above threshold
            self.assertIsNotNone(results_threshold)
        except (ValueError, NotImplementedError):
            # Skip if not supported
            pass

    def test_delete_by_ids(self):
        """Test deleting documents by their IDs"""
        # Create a fresh index for this test to avoid interference
        test_index = f"{self.test_index_name}_delete_by_ids"
        self.test_indexes.add(test_index)
        
        try:
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=test_index,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            # Add test documents
            document_ids = vector_store.add_texts(
                texts=self.test_texts, 
                metadatas=self.test_metadatas
            )
            
            # Get some IDs to delete (first 2 documents)
            ids_to_delete = document_ids[:2]
            
            # Delete the documents
            result = vector_store.delete(ids=ids_to_delete)
            self.assertTrue(result, "Delete operation failed")
            
            # Try to retrieve deleted documents
            retrieved_docs = vector_store.get_by_ids(ids_to_delete)
            self.assertEqual(len(retrieved_docs), 0, 
                            "Deleted documents still retrievable by ID")
            
            # Verify other documents still exist
            remaining_ids = document_ids[2:4]
            retrieved_remaining = vector_store.get_by_ids(remaining_ids)
            self.assertGreater(len(retrieved_remaining), 0,
                            "Non-deleted documents were affected")
        finally:
            try:
                self.nd.delete_index(name=test_index)
            except Exception:
                pass

    def test_delete_by_filter(self):
        """Test deleting documents by filter"""
        # Create a fresh index for this test to avoid interference
        test_index = f"{self.test_index_name}_delete_by_filter"
        self.test_indexes.add(test_index)
        
        try:
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=test_index,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            # Add test documents
            vector_store.add_texts(texts=self.test_texts, metadatas=self.test_metadatas)
            
            # Use filter to delete programming documents (string values only)
            filter_to_delete = [{"category": {"$eq": "programming"}}]
            result = vector_store.delete(filter=filter_to_delete)
            self.assertTrue(result, "Delete by filter operation failed")
            
            # Verify deletion by searching for programming content
            programming_query = "JavaScript programming"
            results_after_delete = vector_store.similarity_search(
                programming_query, k=10
            )
            
            # Check that no simple programming documents remain
            for doc in results_after_delete:
                category = doc.metadata.get("category", "")
                if isinstance(category, str):
                    self.assertNotEqual(category.lower(), "programming",
                                    "Programming document was not deleted")
                # Note: Documents with category as list containing "programming" 
                # will still exist as $eq only matches exact field values, not list contents
        finally:
            try:
                self.nd.delete_index(name=test_index)
            except Exception:
                pass

    def test_delete_by_filter_with_in_operator(self):
        """Test $in operator behavior: matches exact field values only, not list contents"""
        # Create a fresh index for this test to avoid interference
        test_index = f"{self.test_index_name}_delete_in"
        self.test_indexes.add(test_index)
        
        try:
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=test_index,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            # Add documents
            added_ids = vector_store.add_texts(texts=self.test_texts, metadatas=self.test_metadatas)
            # print(len(added_ids))
            
            # Count documents with "programming" before deletion
            all_docs_before = vector_store.similarity_search("programming", k=20)
            string_prog_before = sum(1 for doc in all_docs_before 
                                    if isinstance(doc.metadata.get("category"), str) 
                                    and doc.metadata.get("category").lower() == "programming")
            list_prog_before = sum(1 for doc in all_docs_before 
                                if isinstance(doc.metadata.get("category"), list) 
                                and "programming" in [c.lower() for c in doc.metadata.get("category", [])])
            
            # Delete using $in operator - matches only string values equal to "programming"
            filter_to_delete = [{"category": {"$in": ["programming"]}}]
            result = vector_store.delete(filter=filter_to_delete)
            self.assertTrue(result, "Delete with $in operator failed")
            
            # Verify: string "programming" docs should be deleted, but list docs with "programming" should remain
            all_docs_after = vector_store.similarity_search("programming", k=20)
            # for doc in vector_store.get_by_ids(added_ids):
            #     print(doc.page_content)
            # print(len)
            # for doc in all_docs_after:
            #     print(doc.metadata.get("category"))
            #     print(doc.page_content)
            string_prog_after = sum(1 for doc in all_docs_after 
                                if isinstance(doc.metadata.get("category"), str) 
                                and doc.metadata.get("category").lower() == "programming")
            list_prog_after = sum(1 for doc in all_docs_after 
                                if isinstance(doc.metadata.get("category"), list) 
                                and "programming" in [c.lower() for c in doc.metadata.get("category", [])])
            
            # Assertions
            self.assertEqual(string_prog_after, 0, 
                            "String 'programming' documents should be deleted")
            self.assertEqual(list_prog_after, list_prog_before,
                            "List documents containing 'programming' should NOT be deleted by $in operator")
            
        finally:
            try:
                self.nd.delete_index(name=test_index)
            except Exception:
                pass

    def test_delete_nonexistent_id(self):
        """Test deleting a non-existent document ID"""
        fake_id = "nonexistent-id-12345"
        result = self.vector_store.delete(ids=[fake_id])
        
        # Should return False or handle gracefully
        # (behavior may vary based on implementation)
        self.assertIsNotNone(result)

    def test_delete_with_no_parameters(self):
        """Test that delete raises error when called without ids or filter"""
        with self.assertRaises(ValueError):
            self.vector_store.delete()

    def test_get_by_ids(self):
        """Test retrieving documents by their IDs"""
        # Get a subset of IDs
        ids_to_get = self.document_ids[3:6]
        
        retrieved_docs = self.vector_store.get_by_ids(ids_to_get)
        
        self.assertIsNotNone(retrieved_docs)
        self.assertEqual(len(retrieved_docs), len(ids_to_get),
                        "Number of retrieved documents doesn't match requested IDs")
        
        # Verify retrieved documents have correct IDs
        retrieved_ids = [doc.metadata.get("_id") for doc in retrieved_docs]
        for requested_id in ids_to_get:
            self.assertIn(requested_id, retrieved_ids,
                         f"Requested ID {requested_id} not in retrieved documents")

    def test_get_by_ids_empty_list(self):
        """Test get_by_ids with empty list"""
        retrieved_docs = self.vector_store.get_by_ids([])
        self.assertEqual(len(retrieved_docs), 0)

    def test_get_by_ids_nonexistent(self):
        """Test get_by_ids with non-existent IDs"""
        fake_ids = ["fake-id-1", "fake-id-2"]
        retrieved_docs = self.vector_store.get_by_ids(fake_ids)
        
        # Should return empty list for non-existent IDs
        self.assertEqual(len(retrieved_docs), 0,
                        "Retrieved documents for non-existent IDs")

    def test_reconnect_to_existing_index(self):
        """Test reconnecting to an existing index and verifying data persistence"""
        # Reconnect to the same index using from_existing_index
        reconnected_store = EndeeVectorStore.from_existing_index(
            index_name=self.retriever_test_index,
            embedding=self.embed_model,
            api_token=self.endee_api_token
        )
        
        self.assertIsNotNone(reconnected_store)
        self.assertEqual(reconnected_store.index_name, self.retriever_test_index)
        
        # Search for a known AI document
        results = reconnected_store.similarity_search("What is machine learning?", k=2)
        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0, "Reconnected store returned no results")
        
        found = any("machine learning" in doc.page_content.lower() for doc in results)
        self.assertTrue(found, "Reconnected store did not return expected document")

    def test_reconnect_with_constructor(self):
        """Test reconnecting using regular constructor (without force_recreate)"""
        # Reconnect using constructor
        reconnected_store = EndeeVectorStore(
            embedding=self.embed_model,
            api_token=self.endee_api_token,
            index_name=self.retriever_test_index,
            dimension=self.dimension,
            space_type=self.space_type,
            validate_index_config=True
        )
        
        self.assertIsNotNone(reconnected_store)
        
        # Verify data is still there
        results = reconnected_store.similarity_search("database systems", k=2)
        self.assertGreater(len(results), 0, "Reconnected store has no data")

    def test_reconnect_and_add_more_documents(self):
        """Test reconnecting and adding more documents to existing index"""
        # Reconnect
        reconnected_store = EndeeVectorStore.from_existing_index(
            index_name=self.retriever_test_index,
            embedding=self.embed_model,
            api_token=self.endee_api_token
        )
        
        # Add new documents
        new_texts = [
            "Kubernetes is a container orchestration platform.",
            "Docker provides containerization for applications."
        ]
        new_metadatas = [
            {"category": "devops", "tool": "kubernetes"},
            {"category": "devops", "tool": "docker"}
        ]
        
        new_ids = reconnected_store.add_texts(texts=new_texts, metadatas=new_metadatas)
        self.assertEqual(len(new_ids), len(new_texts))
        
        # Verify new documents are searchable
        results = reconnected_store.similarity_search("container orchestration", k=2)
        self.assertGreater(len(results), 0)
        
        found = any("kubernetes" in doc.page_content.lower() or 
                   "docker" in doc.page_content.lower() 
                   for doc in results)
        self.assertTrue(found, "Newly added documents not found after reconnect")

    def test_index_persistence_after_client_recreation(self):
        """Test that index data persists even after recreating the Endee client"""
        from endee import Endee
        
        # Create a new client
        if self.endee_api_token:
            new_client = Endee(token=self.endee_api_token)
        else:
            new_client = Endee(token="")
        
        # Create vector store with new client
        new_store = EndeeVectorStore(
            embedding=self.embed_model,
            endee_client=new_client,
            index_name=self.retriever_test_index
        )
        
        # Verify original data is accessible
        results = new_store.similarity_search("machine learning", k=2)
        self.assertGreater(len(results), 0, "Data not accessible with new client")
        
        found = any("machine learning" in doc.page_content.lower() for doc in results)
        self.assertTrue(found, "Original data not found with new client")

    def test_multiple_vector_stores_same_index(self):
        """Test that multiple vector store instances can access the same index"""
        # Create second instance pointing to same index
        store2 = EndeeVectorStore.from_existing_index(
            index_name=self.retriever_test_index,
            embedding=self.embed_model,
            api_token=self.endee_api_token
        )
        
        # Both should return similar results
        query = "Python programming"
        results1 = self.vector_store.similarity_search(query, k=2)
        results2 = store2.similarity_search(query, k=2)
        
        self.assertEqual(len(results1), len(results2),
                        "Different instances returned different result counts")
        
        # Results should be identical (same IDs)
        ids1 = sorted([doc.metadata.get("_id") for doc in results1])
        ids2 = sorted([doc.metadata.get("_id") for doc in results2])
        self.assertEqual(ids1, ids2,
                        "Different instances returned different results")

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