#!/usr/bin/env python3
"""
Test Hybrid Search Functionality

This test suite comprehensively tests the hybrid search capabilities of EndeeVectorStore,
which combines dense and sparse embeddings for improved retrieval performance.

Tests include:
- Creating hybrid indexes
- Adding documents with both dense and sparse embeddings
- Performing hybrid searches
- Comparing hybrid vs dense search results
- Using filters with hybrid search
- Factory methods (from_texts, from_documents) in hybrid mode
- Edge cases and error handling
"""

import unittest
import uuid
from typing import List

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings

from langchain_endee import EndeeVectorStore, RetrievalMode, Precision
from langchain_endee.sparse_embeddings import SparseEmbeddings, SparseVector
from tests.setup_class import EndeeLangChainTestSetup


# ==============================================================================
# Mock Embeddings for Testing
# ==============================================================================

class MockDenseEmbeddings(Embeddings):
    """Mock dense embeddings for testing."""

    def __init__(self, dimension: int = 384):
        self.dimension = dimension

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Generate deterministic embeddings based on text hash."""
        embeddings = []
        for text in texts:
            hash_val = hash(text)
            embedding = [(hash_val + i) % 100 / 100.0 for i in range(self.dimension)]
            embeddings.append(embedding)
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        """Generate embedding for query."""
        return self.embed_documents([text])[0]


class MockSparseEmbeddings(SparseEmbeddings):
    """Mock sparse embeddings for testing hybrid search."""

    def __init__(self, sparse_dim: int = 1000):
        self.sparse_dim = sparse_dim

    def embed_documents(self, texts: List[str]) -> List[SparseVector]:
        """Generate mock sparse embeddings."""
        sparse_vectors = []
        for text in texts:
            # Create sparse representation based on words
            words = text.lower().split()
            # Use hash to create deterministic indices
            indices = sorted(list(set([hash(word) % self.sparse_dim for word in words[:20]])))
            # Create weights based on word position (TF-IDF-like)
            values = [1.0 / (i + 1) for i in range(len(indices))]
            sparse_vectors.append(SparseVector(indices=indices, values=values))
        return sparse_vectors

    def embed_query(self, text: str) -> SparseVector:
        """Generate sparse embedding for query."""
        return self.embed_documents([text])[0]

    def get_sparse_dim(self) -> int:
        """Get the sparse dimensionality."""
        return self.sparse_dim


# ==============================================================================
# Test Cases
# ==============================================================================

class TestHybridSearch(EndeeLangChainTestSetup):
    """Test suite for hybrid search functionality."""

    @classmethod
    def setUpClass(cls):
        """Set up test fixtures."""
        super().setUpClass()

        # Initialize embeddings
        cls.dense_embeddings = MockDenseEmbeddings(dimension=cls.dimension)
        cls.sparse_embeddings = MockSparseEmbeddings(sparse_dim=1000)

        # Hybrid test data
        cls.hybrid_texts = [
            "Python is a high-level programming language with dynamic typing",
            "Machine learning algorithms learn patterns from data",
            "Vector databases store and retrieve high-dimensional embeddings",
            "Natural language processing enables computers to understand text",
            "Deep learning uses neural networks with multiple layers",
            "Information retrieval finds relevant documents for queries",
            "Hybrid search combines lexical and semantic search methods",
            "BM25 is a probabilistic ranking function for text retrieval",
            "Dense embeddings capture semantic meaning in continuous space",
            "Sparse embeddings represent text with discrete vocabulary indices",
        ]

        cls.hybrid_metadatas = [
            {"category": "programming", "topic": "python", "difficulty": "beginner"},
            {"category": "ai", "topic": "machine_learning", "difficulty": "intermediate"},
            {"category": "database", "topic": "vector_db", "difficulty": "intermediate"},
            {"category": "ai", "topic": "nlp", "difficulty": "intermediate"},
            {"category": "ai", "topic": "deep_learning", "difficulty": "advanced"},
            {"category": "database", "topic": "information_retrieval", "difficulty": "intermediate"},
            {"category": "database", "topic": "hybrid_search", "difficulty": "advanced"},
            {"category": "database", "topic": "bm25", "difficulty": "advanced"},
            {"category": "ai", "topic": "embeddings", "difficulty": "intermediate"},
            {"category": "ai", "topic": "embeddings", "difficulty": "intermediate"},
        ]

    def test_01_create_hybrid_index(self):
        """Test creating a vector store with hybrid mode."""
        print("\n=== Test 01: Create Hybrid Index ===")

        index_name = f"{self.test_index_name}_hybrid_01"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            precision=Precision.INT8D,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        self.assertIsNotNone(vector_store)
        self.assertEqual(vector_store.retrieval_mode, RetrievalMode.HYBRID)
        self.assertIsNotNone(vector_store.sparse_embeddings)
        print("✅ Hybrid index created successfully")

    def test_02_add_texts_hybrid(self):
        """Test adding texts to hybrid index."""
        print("\n=== Test 02: Add Texts to Hybrid Index ===")

        index_name = f"{self.test_index_name}_hybrid_02"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        # Add texts
        texts = self.hybrid_texts[:5]
        metadatas = self.hybrid_metadatas[:5]

        ids = vector_store.add_texts(
            texts=texts,
            metadatas=metadatas,
        )

        self.assertEqual(len(ids), len(texts))
        print(f"✅ Added {len(ids)} texts to hybrid index")

    def test_03_hybrid_similarity_search(self):
        """Test hybrid similarity search."""
        print("\n=== Test 03: Hybrid Similarity Search ===")

        index_name = f"{self.test_index_name}_hybrid_03"
        self.test_indexes.add(index_name)

        # Create and populate hybrid index
        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts,
            metadatas=self.hybrid_metadatas,
        )

        # Perform hybrid search
        query = "machine learning and neural networks"
        results = vector_store.similarity_search(query, k=3)

        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)
        self.assertLessEqual(len(results), 3)

        print(f"✅ Hybrid search returned {len(results)} results")
        print(f"   Query: '{query}'")
        for i, doc in enumerate(results, 1):
            print(f"   {i}. {doc.page_content[:60]}...")

    def test_04_hybrid_search_with_scores(self):
        """Test hybrid search with similarity scores."""
        print("\n=== Test 04: Hybrid Search with Scores ===")

        index_name = f"{self.test_index_name}_hybrid_04"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts,
            metadatas=self.hybrid_metadatas,
        )

        # Search with scores
        query = "vector database embeddings"
        results = vector_store.similarity_search_with_score(query, k=5)

        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

        # Verify all results have scores
        for doc, score in results:
            self.assertIsInstance(doc, Document)
            self.assertIsInstance(score, float)
            self.assertGreaterEqual(score, 0.0)

        print(f"✅ Hybrid search with scores returned {len(results)} results")
        print(f"   Query: '{query}'")
        for i, (doc, score) in enumerate(results, 1):
            print(f"   {i}. [Score: {score:.4f}] {doc.page_content[:50]}...")

    def test_05_hybrid_search_with_filter(self):
        """Test hybrid search with metadata filtering."""
        print("\n=== Test 05: Hybrid Search with Filter ===")

        index_name = f"{self.test_index_name}_hybrid_05"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts,
            metadatas=self.hybrid_metadatas,
        )

        # Search with filter
        query = "learning algorithms"
        filter_condition = [{"category": {"$eq": "ai"}}]

        results = vector_store.similarity_search(
            query,
            k=5,
            filter=filter_condition,
        )

        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

        # Verify all results match filter
        for doc in results:
            self.assertEqual(doc.metadata.get("category"), "ai")

        print(f"✅ Filtered hybrid search returned {len(results)} results")
        print(f"   Query: '{query}'")
        print(f"   Filter: {filter_condition}")
        for i, doc in enumerate(results, 1):
            print(f"   {i}. {doc.page_content[:50]}... (category: {doc.metadata.get('category')})")

    def test_06_hybrid_vs_dense_comparison(self):
        """Compare hybrid search results with dense-only search."""
        print("\n=== Test 06: Hybrid vs Dense Search Comparison ===")

        # Create dense index
        dense_index_name = f"{self.test_index_name}_dense_06"
        self.test_indexes.add(dense_index_name)

        dense_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=dense_index_name,
            embedding=self.dense_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.DENSE,
            force_recreate=True,
        )

        # Create hybrid index
        hybrid_index_name = f"{self.test_index_name}_hybrid_06"
        self.test_indexes.add(hybrid_index_name)

        hybrid_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=hybrid_index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        # Add same texts to both
        texts = self.hybrid_texts[:7]
        metadatas = self.hybrid_metadatas[:7]

        dense_store.add_texts(texts=texts, metadatas=metadatas)
        hybrid_store.add_texts(texts=texts, metadatas=metadatas)

        # Compare search results
        query = "BM25 probabilistic ranking"

        dense_results = dense_store.similarity_search_with_score(query, k=3)
        hybrid_results = hybrid_store.similarity_search_with_score(query, k=3)

        self.assertIsNotNone(dense_results)
        self.assertIsNotNone(hybrid_results)

        print(f"✅ Comparison complete")
        print(f"   Query: '{query}'")
        print(f"\n   Dense search results:")
        for i, (doc, score) in enumerate(dense_results, 1):
            print(f"   {i}. [Score: {score:.4f}] {doc.page_content[:50]}...")

        print(f"\n   Hybrid search results:")
        for i, (doc, score) in enumerate(hybrid_results, 1):
            print(f"   {i}. [Score: {score:.4f}] {doc.page_content[:50]}...")

    def test_07_from_texts_hybrid(self):
        """Test from_texts factory method with hybrid mode."""
        print("\n=== Test 07: from_texts with Hybrid Mode ===")

        index_name = f"{self.test_index_name}_hybrid_07"
        self.test_indexes.add(index_name)

        texts = self.hybrid_texts[:5]
        metadatas = self.hybrid_metadatas[:5]

        vector_store = EndeeVectorStore.from_texts(
            texts=texts,
            embedding=self.dense_embeddings,
            metadatas=metadatas,
            api_token=self.endee_api_token,
            index_name=index_name,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            sparse_embedding=self.sparse_embeddings,
            force_recreate=True,
        )

        self.assertIsNotNone(vector_store)
        self.assertEqual(vector_store.retrieval_mode, RetrievalMode.HYBRID)

        # Test search
        results = vector_store.similarity_search("programming language", k=2)
        self.assertGreater(len(results), 0)

        print(f"✅ from_texts created hybrid index with {len(texts)} texts")
        print(f"   Search returned {len(results)} results")

    def test_08_from_documents_hybrid(self):
        """Test from_documents factory method with hybrid mode."""
        print("\n=== Test 08: from_documents with Hybrid Mode ===")

        index_name = f"{self.test_index_name}_hybrid_08"
        self.test_indexes.add(index_name)

        documents = [
            Document(
                page_content=text,
                metadata=meta
            )
            for text, meta in zip(self.hybrid_texts[:5], self.hybrid_metadatas[:5])
        ]

        vector_store = EndeeVectorStore.from_documents(
            documents=documents,
            embedding=self.dense_embeddings,
            api_token=self.endee_api_token,
            index_name=index_name,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            sparse_embedding=self.sparse_embeddings,
            force_recreate=True,
        )

        self.assertIsNotNone(vector_store)
        self.assertEqual(vector_store.retrieval_mode, RetrievalMode.HYBRID)

        # Test search with filter
        results = vector_store.similarity_search(
            "artificial intelligence",
            k=3,
            filter=[{"category": {"$eq": "ai"}}]
        )

        print(f"✅ from_documents created hybrid index with {len(documents)} documents")
        print(f"   Filtered search returned {len(results)} results")

    def test_09_hybrid_search_by_vector(self):
        """Test hybrid search using pre-computed vectors."""
        print("\n=== Test 09: Hybrid Search by Vector ===")

        index_name = f"{self.test_index_name}_hybrid_09"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts[:6],
            metadatas=self.hybrid_metadatas[:6],
        )

        # Generate query vectors
        query_text = "deep learning neural networks"
        dense_vector = self.dense_embeddings.embed_query(query_text)
        sparse_vector = self.sparse_embeddings.embed_query(query_text)

        # Search by vector
        results = vector_store.similarity_search_by_vector_with_score(
            embedding=dense_vector,
            k=3,
            sparse_indices=sparse_vector.indices,
            sparse_values=sparse_vector.values,
        )

        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)

        print(f"✅ Hybrid search by vector returned {len(results)} results")
        for i, (doc, score) in enumerate(results, 1):
            print(f"   {i}. [Score: {score:.4f}] {doc.page_content[:50]}...")

    def test_10_hybrid_as_retriever(self):
        """Test using hybrid vector store as a retriever."""
        print("\n=== Test 10: Hybrid Store as Retriever ===")

        index_name = f"{self.test_index_name}_hybrid_10"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts,
            metadatas=self.hybrid_metadatas,
        )

        # Create retriever
        retriever = vector_store.as_retriever(
            search_kwargs={"k": 3}
        )

        # Use retriever
        query = "information retrieval systems"
        results = retriever.invoke(query)

        self.assertIsNotNone(results)
        self.assertGreater(len(results), 0)
        self.assertLessEqual(len(results), 3)

        print(f"✅ Hybrid retriever returned {len(results)} results")
        print(f"   Query: '{query}'")
        for i, doc in enumerate(results, 1):
            print(f"   {i}. {doc.page_content[:50]}...")

    def test_11_hybrid_batch_operations(self):
        """Test batch operations with hybrid search."""
        print("\n=== Test 11: Hybrid Batch Operations ===")

        index_name = f"{self.test_index_name}_hybrid_11"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        # Add texts in batches
        batch_size = 3
        all_ids = []

        for i in range(0, len(self.hybrid_texts), batch_size):
            batch_texts = self.hybrid_texts[i:i+batch_size]
            batch_metadatas = self.hybrid_metadatas[i:i+batch_size]

            ids = vector_store.add_texts(
                texts=batch_texts,
                metadatas=batch_metadatas,
                batch_size=batch_size,
            )
            all_ids.extend(ids)

        self.assertEqual(len(all_ids), len(self.hybrid_texts))

        # Test search
        results = vector_store.similarity_search("search algorithms", k=3)
        self.assertGreater(len(results), 0)

        print(f"✅ Added {len(all_ids)} texts in batches of {batch_size}")
        print(f"   Search returned {len(results)} results")

    def test_12_hybrid_delete_and_verify(self):
        """Test deleting documents from hybrid index."""
        print("\n=== Test 12: Delete from Hybrid Index ===")

        index_name = f"{self.test_index_name}_hybrid_12"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        # Add texts
        texts = self.hybrid_texts[:5]
        metadatas = self.hybrid_metadatas[:5]
        ids = vector_store.add_texts(texts=texts, metadatas=metadatas)

        # Delete one document
        delete_id = ids[0]
        result = vector_store.delete(ids=[delete_id])

        self.assertTrue(result)

        # Verify deletion
        retrieved = vector_store.get_by_ids([delete_id])
        self.assertEqual(len(retrieved), 0)

        # Verify remaining documents
        remaining = vector_store.get_by_ids(ids[1:])
        self.assertEqual(len(remaining), len(ids) - 1)

        print(f"✅ Deleted document with ID: {delete_id}")
        print(f"   Verified {len(remaining)} documents remain")

    def test_13_hybrid_multiple_filters(self):
        """Test hybrid search with multiple filter conditions."""
        print("\n=== Test 13: Hybrid Search with Multiple Filters ===")

        index_name = f"{self.test_index_name}_hybrid_13"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts,
            metadatas=self.hybrid_metadatas,
        )

        # Search with multiple filters
        query = "advanced concepts"
        filter_conditions = [
            {"category": {"$eq": "ai"}},
            {"difficulty": {"$eq": "advanced"}}
        ]

        results = vector_store.similarity_search(
            query,
            k=5,
            filter=filter_conditions,
        )

        self.assertIsNotNone(results)

        # Verify all results match both filters
        for doc in results:
            self.assertEqual(doc.metadata.get("category"), "ai")
            self.assertEqual(doc.metadata.get("difficulty"), "advanced")

        print(f"✅ Multi-filter hybrid search returned {len(results)} results")
        print(f"   All results match: category='ai' AND difficulty='advanced'")

    def test_14_hybrid_empty_query(self):
        """Test hybrid search behavior with edge cases."""
        print("\n=== Test 14: Hybrid Search Edge Cases ===")

        index_name = f"{self.test_index_name}_hybrid_14"
        self.test_indexes.add(index_name)

        vector_store = EndeeVectorStore(
            api_token=self.endee_api_token,
            index_name=index_name,
            embedding=self.dense_embeddings,
            sparse_embedding=self.sparse_embeddings,
            dimension=self.dimension,
            space_type=self.space_type,
            retrieval_mode=RetrievalMode.HYBRID,
            force_recreate=True,
        )

        vector_store.add_texts(
            texts=self.hybrid_texts[:5],
            metadatas=self.hybrid_metadatas[:5],
        )

        # Test with very short query
        short_query = "ML"
        results = vector_store.similarity_search(short_query, k=2)
        self.assertIsNotNone(results)

        # Test with single word
        single_word = "embeddings"
        results = vector_store.similarity_search(single_word, k=2)
        self.assertIsNotNone(results)

        # Test with very long query
        long_query = " ".join(self.hybrid_texts[:3])
        results = vector_store.similarity_search(long_query, k=2)
        self.assertIsNotNone(results)

        print(f"✅ Edge case tests passed:")
        print(f"   - Short query: '{short_query}'")
        print(f"   - Single word: '{single_word}'")
        print(f"   - Long query: {len(long_query)} characters")


# ==============================================================================
# Test Runner
# ==============================================================================

if __name__ == '__main__':
    # Run tests with verbose output
    unittest.main(verbosity=2)
