import unittest
import sys
from pathlib import Path
from tests.setup_class import EndeeLangChainTestSetup

sys.path.append(str(Path(__file__).parent.parent))

from langchain_endee.vectorstores import EndeeVectorStore
from langchain_community.embeddings import HuggingFaceEmbeddings 
from endee import Precision

class TestEndeeVectorStore(EndeeLangChainTestSetup):
    def setUp(self):
        self.embed_model = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={"device": "cpu"}
        )

    def test_create_vector_store_basic(self):
        """Test creating Endee vector store with basic parameters"""
        try:
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=self.test_index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            self.assertIsNotNone(vector_store)
            self.assertEqual(vector_store.index_name, self.test_index_name)
            
            # Verify index exists in Endee
            index = self.nd.get_index(name=self.test_index_name)
            self.assertIsNotNone(index)
            index_info = index.describe()
            self.assertEqual(index_info["dimension"], self.dimension)
            self.assertEqual(index_info["space_type"], self.space_type)
        except Exception as e:
            self.fail(f"Vector store creation failed: {str(e)}")

    def test_create_vector_store_different_precision(self):
        """Test creating Endee vector store with different precision levels"""
        # Updated precision types from the new implementation
        precisions = [
            Precision.BINARY2,
            Precision.FLOAT16, 
            Precision.INT8D,
            Precision.INT16D,
            Precision.FLOAT32
        ]

        for precision in precisions:
            with self.subTest(precision=precision):
                precision_name = str(precision).replace(".", "_")
                index_name = f"{self.test_index_name}_{precision_name}"
                try:
                    # Create vector store with precision
                    vector_store = EndeeVectorStore(
                        embedding=self.embed_model,
                        api_token=self.endee_api_token,
                        index_name=index_name,
                        dimension=self.dimension,
                        space_type=self.space_type,
                        precision=precision,
                        force_recreate=True
                    )

                    self.assertIsNotNone(vector_store)

                    # Index should exist
                    index = self.nd.get_index(name=index_name)
                    self.assertIsNotNone(index)

                    info = index.describe()
                    self.assertEqual(info["dimension"], self.dimension)
                    self.assertEqual(info["space_type"], self.space_type)
                    self.assertEqual(info["precision"], precision)

                except Exception as e:
                    self.fail(f"Precision '{precision}' failed: {str(e)}")

                finally:
                    # ALWAYS delete the index to avoid conflicts
                    try:
                        self.nd.delete_index(name=index_name)
                    except Exception:
                        # ignore if already deleted or not created
                        pass

    def test_create_vector_store_local_deployment(self):
        """Test creating vector store without API token (local deployment)"""
        index_name = f"{self.test_index_name}_local"
        self.test_indexes.add(index_name)
        
        try:
            # Create vector store without API token (should work with local deployment)
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=None,  # No API token for local deployment
                index_name=index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            self.assertIsNotNone(vector_store)
            self.assertEqual(vector_store.index_name, index_name)
            
            # Add documents
            ids = vector_store.add_texts(texts=self.test_texts, metadatas=self.test_metadatas)
            self.assertEqual(len(ids), len(self.test_texts))
            
            # Verify search works
            test_query = "Python programming"
            results = vector_store.similarity_search(test_query, k=2)
            self.assertGreater(len(results), 0, "No results found in index")
            
        except Exception as e:
            # This test may fail if local Endee server is not running
            if "not running" in str(e).lower():
                self.skipTest("Local Endee server not running on http://127.0.0.1:8080")
            else:
                self.fail(f"Vector store creation without API token failed: {str(e)}")
        finally:
            try:
                self.nd.delete_index(name=index_name)
            except Exception:
                pass

    def test_create_vector_store_with_documents(self):
        """Test creating vector store and adding documents"""
        try:
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=self.test_index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            ids = vector_store.add_texts(texts=self.test_texts, metadatas=self.test_metadatas)
            self.assertEqual(len(ids), len(self.test_texts))
            
            # Verify documents were added by performing a search
            test_query = "Python programming"
            results = vector_store.similarity_search(test_query, k=len(self.test_texts))
            self.assertGreater(len(results), 0, "No vectors found in index after adding documents")
        except Exception as e:
            self.fail(f"Vector store creation with documents failed: {str(e)}")

    def test_from_texts_method(self):
        """Test creating vector store using from_texts class method"""
        index_name = f"{self.test_index_name}_from_texts"
        self.test_indexes.add(index_name)
        
        try:
            # Create vector store from texts
            vector_store = EndeeVectorStore.from_texts(
                texts=self.test_texts,
                embedding=self.embed_model,
                metadatas=self.test_metadatas,
                api_token=self.endee_api_token,
                index_name=index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            self.assertIsNotNone(vector_store)
            
            # Verify documents were added
            test_query = "Python programming"
            results = vector_store.similarity_search(test_query, k=2)
            self.assertGreater(len(results), 0)
            
        except Exception as e:
            self.fail(f"from_texts method failed: {str(e)}")
        finally:
            try:
                self.nd.delete_index(name=index_name)
            except Exception:
                pass

    def test_from_documents_method(self):
        """Test creating vector store using from_documents class method"""
        from langchain_core.documents import Document
        
        index_name = f"{self.test_index_name}_from_docs"
        self.test_indexes.add(index_name)
        
        try:
            # Create Document objects
            documents = [
                Document(page_content=text, metadata=metadata)
                for text, metadata in zip(self.test_texts, self.test_metadatas)
            ]
            
            # Create vector store from documents
            vector_store = EndeeVectorStore.from_documents(
                documents=documents,
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            self.assertIsNotNone(vector_store)
            
            # Verify documents were added
            test_query = "Python programming"
            results = vector_store.similarity_search(test_query, k=2)
            self.assertGreater(len(results), 0)
            
        except Exception as e:
            self.fail(f"from_documents method failed: {str(e)}")
        finally:
            try:
                self.nd.delete_index(name=index_name)
            except Exception:
                pass

    def test_from_existing_index(self):
        """Test retrieving an existing index using from_existing_index method"""
        index_name = f"{self.test_index_name}_existing"
        self.test_indexes.add(index_name)
        
        try:
            # Create index first
            self.nd.create_index(
                name=index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            # Now retrieve it using from_existing_index
            vector_store = EndeeVectorStore.from_existing_index(
                index_name=index_name,
                embedding=self.embed_model,
                api_token=self.endee_api_token
            )
            
            self.assertIsNotNone(vector_store)
            self.assertEqual(vector_store.index_name, index_name)
            
        except Exception as e:
            self.fail(f"from_existing_index method failed: {str(e)}")
        finally:
            try:
                self.nd.delete_index(name=index_name)
            except Exception:
                pass

    def test_force_recreate_index(self):
        """Test force_recreate parameter to delete and recreate existing index"""
        index_name = f"{self.test_index_name}_recreate"
        self.test_indexes.add(index_name)
        
        try:
            # Create index first
            vector_store1 = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            # Add some documents
            ids1 = vector_store1.add_texts(texts=self.test_texts[:3])
            self.assertEqual(len(ids1), 3)
            
            # Recreate with force_recreate=True
            vector_store2 = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=index_name,
                dimension=self.dimension,
                space_type=self.space_type,
                force_recreate=True
            )
            
            # Index should be empty now
            results = vector_store2.similarity_search("Python", k=10)
            self.assertEqual(len(results), 0, "Index should be empty after force_recreate")
            
        except Exception as e:
            self.fail(f"force_recreate test failed: {str(e)}")
        finally:
            try:
                self.nd.delete_index(name=index_name)
            except Exception:
                pass

    def test_create_vector_store_invalid_params(self):
        """Test vector store creation with invalid parameters"""
        # Invalid dimension
        unique_index_name = f"{self.test_index_name}_invalid_dim"
        self.test_indexes.add(unique_index_name)
        with self.assertRaises(Exception):
            EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=unique_index_name,
                dimension=-1,
                space_type=self.space_type
            )
        
        # Invalid space type
        unique_index_name = f"{self.test_index_name}_invalid_space"
        self.test_indexes.add(unique_index_name)
        with self.assertRaises(Exception):
            EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=unique_index_name,
                dimension=self.dimension,
                space_type="invalid_space"
            )
        
        # Invalid API token (only if not using local deployment)
        if self.endee_api_token:
            unique_index_name = f"{self.test_index_name}_invalid_token"
            self.test_indexes.add(unique_index_name)
            with self.assertRaises(Exception):
                EndeeVectorStore(
                    embedding=self.embed_model,
                    api_token="invalid:invalid_key:dev",
                    index_name=unique_index_name,
                    dimension=self.dimension,
                    space_type=self.space_type
                )
        
        # Missing required parameters - index_name
        with self.assertRaises(ValueError):
            EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                dimension=self.dimension,
                space_type=self.space_type
            )
        
        # Missing required parameters - embedding
        with self.assertRaises(ValueError):
            EndeeVectorStore(
                embedding=None,
                api_token=self.endee_api_token,
                index_name=self.test_index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )

    def test_auto_detect_embedding_model_type(self):
        """Test automatic detection of embedding model type"""
        try:
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=self.test_index_name,
                dimension=self.dimension,
                space_type=self.space_type
            )
            
            # Should auto-detect HuggingFace embeddings
            self.assertEqual(vector_store.embedding_model_type, "huggingface")
            # Should set appropriate max_text_length
            self.assertEqual(vector_store.max_text_length, 512)
            
        except Exception as e:
            self.fail(f"Auto-detect embedding model type failed: {str(e)}")

    def test_custom_max_text_length(self):
        """Test setting custom max_text_length"""
        index_name = f"{self.test_index_name}_custom_length"
        self.test_indexes.add(index_name)
        
        try:
            custom_length = 1024
            vector_store = EndeeVectorStore(
                embedding=self.embed_model,
                api_token=self.endee_api_token,
                index_name=index_name,
                dimension=self.dimension,
                space_type=self.space_type,
                max_text_length=custom_length
            )
            
            self.assertEqual(vector_store.max_text_length, custom_length)
            
        except Exception as e:
            self.fail(f"Custom max_text_length test failed: {str(e)}")
        finally:
            try:
                self.nd.delete_index(name=index_name)
            except Exception:
                pass

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