import os
import time
from typing import Any

import pytest
from langchain_community.embeddings.fake import DeterministicFakeEmbedding
from lfx.components.mongodb import MongoVectorStoreComponent
from lfx.schema.data import Data
from pymongo.collection import Collection

from tests.base import ComponentTestBaseWithoutClient, VersionComponentMapping


@pytest.mark.skipif(
    not os.environ.get("MONGODB_ATLAS_URI"), reason="Environment variable MONGODB_ATLAS_URI is not defined."
)
class TestMongoVectorStoreComponent(ComponentTestBaseWithoutClient):
    @pytest.fixture
    def component_class(self) -> type[Any]:
        """Return the component class to test."""
        return MongoVectorStoreComponent

    @pytest.fixture
    def default_kwargs(self) -> dict[str, Any]:
        """Return the default kwargs for the component."""
        return {
            "mongodb_atlas_cluster_uri": os.getenv("MONGODB_ATLAS_URI"),
            "db_name": "test_db",
            "collection_name": "test_collection",
            "index_name": "test_index",
            "enable_mtls": False,
            "embedding": DeterministicFakeEmbedding(size=8),
            "index_field": "embedding",
            "filter_field": "text",
            "number_dimensions": 8,
            "similarity": "cosine",
            "quantization": "scalar",
            "insert_mode": "append",
            "ingest_data": [Data(data={"text": "test data 1"}), Data(data={"text": "test data 2"})],
        }

    @pytest.fixture
    def file_names_mapping(self) -> list[VersionComponentMapping]:
        """Return the file names mapping for different versions."""
        return [
            {"version": "1.0.19", "module": "vectorstores", "file_name": "MongoDBAtlasVector"},
            {"version": "1.1.0", "module": "vectorstores", "file_name": "mongodb_atlas"},
            {"version": "1.1.1", "module": "vectorstores", "file_name": "mongodb_atlas"},
        ]

    def __create_search_index(
        self, component_class: type[MongoVectorStoreComponent], collection: Collection, default_kwargs: dict[str, Any]
    ) -> None:
        """Create a vector search index if it doesn't exist."""
        component_class().set(**default_kwargs).verify_search_index(collection)

        # Verify index was created
        indexes = collection.list_search_indexes()
        index_names = {idx["name"]: idx["type"] for idx in indexes}
        index_type = index_names.get(default_kwargs["index_name"])
        assert default_kwargs["index_name"] in index_names
        assert index_type == "vectorSearch"

    def test_create_db(self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]) -> None:
        """Test creating a MongoDB Atlas vector store."""
        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)
        vector_store = component.build_vector_store()
        assert vector_store is not None
        # Access MongoDB collection through the vector store's internal client
        assert vector_store._collection.name == default_kwargs["collection_name"]
        assert vector_store._index_name == default_kwargs["index_name"]

    def test_create_collection_with_data(
        self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test creating a collection with data."""
        test_texts = ["test data 1", "test data 2", "something completely different"]
        default_kwargs["ingest_data"] = [Data(data={"text": text}) for text in test_texts]
        default_kwargs["insert_mode"] = "overwrite"

        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)
        vector_store = component.build_vector_store()

        # Verify collection exists and has the correct data
        collection = vector_store._collection
        assert collection.name == default_kwargs["collection_name"]
        assert collection.count_documents({}) == len(test_texts)

    def test_similarity_search(
        self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test the similarity search functionality."""
        # Create test data with distinct topics
        test_data = [
            "The quick brown fox jumps over the lazy dog",
            "Python is a popular programming language",
            "Machine learning models process data",
            "The lazy dog sleeps all day long",
        ]
        default_kwargs["ingest_data"] = [Data(data={"text": text, "metadata": {}}) for text in test_data]
        default_kwargs["number_of_results"] = 2
        default_kwargs["insert_mode"] = "overwrite"

        # Create and initialize the component
        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)

        # Build the vector store first to ensure data is ingested
        vector_store = component.build_vector_store()
        assert vector_store is not None

        # Verify documents were stored with embeddings
        documents = list(vector_store._collection.find({}))
        assert len(documents) == len(test_data)
        for doc in documents:
            assert "embedding" in doc
            assert isinstance(doc["embedding"], list)
            assert len(doc["embedding"]) == 8  # Should match our embedding size

        self.__create_search_index(component_class, vector_store._collection, default_kwargs)

        # Verify index was created
        indexes = vector_store._collection.list_search_indexes()
        index_names = [idx["name"] for idx in indexes]
        assert default_kwargs["index_name"] in index_names

        # Test similarity search through the component
        component.set(search_query="dog")
        results = component.search_documents()
        time.sleep(5)  # wait the results come from API

        assert len(results) == 2, "Expected 2 results for 'lazy dog' query"
        # The most relevant results should be about dogs
        assert any("dog" in result.data["text"].lower() for result in results)

        # Test with different number of results
        component.set(number_of_results=3)
        results = component.search_documents()
        assert len(results) == 3
        assert all("text" in result.data for result in results)

    def test_mtls_configuration(
        self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test mTLS configuration handling."""
        # Test with invalid mTLS configuration
        default_kwargs["enable_mtls"] = True
        default_kwargs["mongodb_atlas_client_cert"] = "invalid-cert-content"

        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)
        with pytest.raises(ValueError, match="Failed to connect to MongoDB Atlas"):
            component.build_vector_store()

    def test_empty_search_query(
        self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test search with empty query."""
        default_kwargs["insert_mode"] = "overwrite"
        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)

        # Test with empty search query
        component.set(search_query="")
        results = component.search_documents()
        assert len(results) == 0

    def test_metadata_handling(
        self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test handling of document metadata."""
        # Create test data with metadata
        test_data = [
            Data(data={"text": "Document 1", "metadata": {"category": "test", "priority": 1}}),
            Data(data={"text": "Document 2", "metadata": {"category": "test", "priority": 2}}),
        ]
        default_kwargs["ingest_data"] = test_data
        default_kwargs["collection_name"] = "test_collection_metadata"

        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)
        vector_store = component.build_vector_store()

        self.__create_search_index(component_class, vector_store._collection, default_kwargs)

        # Test search and verify metadata is preserved
        component.set(search_query="Document", number_of_results=2)
        results = component.search_documents()

        assert len(results) == 2
        for result in results:
            assert "category" in result.data["metadata"]
            assert result.data["metadata"]["category"] == "test"
            assert "priority" in result.data["metadata"]
            assert isinstance(result.data["metadata"]["priority"], int)

    def test_error_handling(
        self, component_class: type[MongoVectorStoreComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test error handling for invalid configurations."""
        component: MongoVectorStoreComponent = component_class().set(**default_kwargs)

        # Test with non-existent database
        default_kwargs["mongodb_atlas_cluster_uri"] = os.getenv("MONGODB_ATLAS_URI")
        default_kwargs["db_name"] = "nonexistent_db"
        component = component_class().set(**default_kwargs)

        # This should not raise an error as MongoDB creates databases and collections on demand
        vector_store = component.build_vector_store()
        assert vector_store is not None
