import os
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch

import pytest
from langflow.services.cache.utils import CACHE_DIR
from lfx.components.vectorstores.local_db import LocalDBComponent
from lfx.schema.data import Data

from tests.base import ComponentTestBaseWithoutClient, VersionComponentMapping


@pytest.mark.api_key_required
class TestLocalDBComponent(ComponentTestBaseWithoutClient):
    @pytest.fixture
    def component_class(self) -> type[Any]:
        """Return the component class to test."""
        return LocalDBComponent

    @pytest.fixture
    def default_kwargs(self, tmp_path: Path) -> dict[str, Any]:
        """Return the default kwargs for the component."""
        from lfx.components.openai.openai import OpenAIEmbeddingsComponent

        from tests.api_keys import get_openai_api_key

        try:
            api_key = get_openai_api_key()
        except ValueError:
            pytest.skip("OPENAI_API_KEY is not set")

        return {
            "embedding": OpenAIEmbeddingsComponent(openai_api_key=api_key).build_embeddings(),
            "collection_name": "test_collection",
            "persist": True,
            "persist_directory": str(tmp_path),  # Convert Path to string
            "mode": "Ingest",
        }

    @pytest.fixture
    def file_names_mapping(self) -> list[VersionComponentMapping]:
        """Return the file names mapping for different versions."""
        # Return an empty list since this is a new component
        return []

    def test_create_db(self, component_class: type[LocalDBComponent], default_kwargs: dict[str, Any]) -> None:
        """Test creating a vector store."""
        component: LocalDBComponent = component_class().set(**default_kwargs)
        component.build_vector_store()
        persist_directory = Path(default_kwargs["persist_directory"])
        assert persist_directory.exists()
        assert persist_directory.is_dir()
        # Assert it isn't empty
        assert len(list(persist_directory.iterdir())) > 0
        # Assert there's a chroma.sqlite3 file (since LocalDB uses Chroma underneath)
        assert (persist_directory / "chroma.sqlite3").exists()
        assert (persist_directory / "chroma.sqlite3").is_file()

    @patch("langchain_chroma.Chroma._collection")
    def test_create_db_with_data(
        self,
        mock_collection,
        component_class: type[LocalDBComponent],
        default_kwargs: dict[str, Any],
    ) -> None:
        """Test creating a vector store with data."""
        # Set ingest_data in default_kwargs to a list of Data objects
        test_texts = ["test data 1", "test data 2", "something completely different"]
        default_kwargs["ingest_data"] = [Data(text=text) for text in test_texts]

        # Mock the collection count to return the expected number
        mock_collection.count.return_value = len(test_texts)
        mock_collection.name = default_kwargs["collection_name"]

        # Mock the _add_documents_to_vector_store method to ensure add_documents is called
        with patch.object(LocalDBComponent, "_add_documents_to_vector_store") as mock_add_docs_method:
            component: LocalDBComponent = component_class().set(**default_kwargs)
            vector_store = component.build_vector_store()

            # Verify the method was called
            mock_add_docs_method.assert_called_once()

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

    def test_default_persist_dir(self, component_class: type[LocalDBComponent], default_kwargs: dict[str, Any]) -> None:
        """Test the default persist directory functionality."""
        # Remove persist_directory from default_kwargs to test default directory
        default_kwargs.pop("persist_directory")

        component: LocalDBComponent = component_class().set(**default_kwargs)

        # Call get_default_persist_dir and check the result
        default_dir = component.get_default_persist_dir()
        expected_dir = Path(CACHE_DIR) / "vector_stores" / default_kwargs["collection_name"]

        assert Path(default_dir) == expected_dir
        assert Path(default_dir).exists()

    @patch("langchain_chroma.Chroma.similarity_search")
    def test_similarity_search(
        self,
        mock_similarity_search,
        component_class: type[LocalDBComponent],
        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(text=text) for text in test_data]
        default_kwargs["search_type"] = "Similarity"
        default_kwargs["number_of_results"] = 2

        # Mock the similarity_search to return documents
        from langchain_core.documents import Document

        mock_docs = [
            Document(page_content="The lazy dog sleeps all day long"),
            Document(page_content="The quick brown fox jumps over the lazy dog"),
        ]
        mock_similarity_search.return_value = mock_docs

        component: LocalDBComponent = component_class().set(**default_kwargs)
        component.build_vector_store()

        # Switch to Retrieve mode
        component.set(mode="Retrieve", search_query="dog sleeping")
        results = component.search_documents()

        assert len(results) == 2
        # The most relevant results should be about dogs
        assert any("dog" in result.text.lower() for result in results)
        mock_similarity_search.assert_called_once_with(query="dog sleeping", k=2)

        # Test with different number of results
        component.set(number_of_results=3)
        another_doc = Document(page_content="Another document")
        mock_similarity_search.return_value = [*mock_docs, another_doc]  # Use unpacking instead of concatenation
        results = component.search_documents()
        assert len(results) == 3

    @patch("langchain_chroma.Chroma.max_marginal_relevance_search")
    def test_mmr_search(
        self,
        mock_mmr_search,
        component_class: type[LocalDBComponent],
        default_kwargs: dict[str, Any],
    ) -> None:
        """Test the MMR search functionality."""
        # Create test data with some similar documents
        test_data = [
            "The quick brown fox jumps",
            "The quick brown fox leaps",
            "The quick brown fox hops",
            "Something completely different about cats",
        ]
        default_kwargs["ingest_data"] = [Data(text=text) for text in test_data]
        default_kwargs["search_type"] = "MMR"
        default_kwargs["number_of_results"] = 3

        # Mock the MMR search to return documents
        from langchain_core.documents import Document

        mock_docs = [
            Document(page_content="The quick brown fox jumps"),
            Document(page_content="The quick brown fox leaps"),
            Document(page_content="Something completely different about cats"),
        ]
        mock_mmr_search.return_value = mock_docs

        component: LocalDBComponent = component_class().set(**default_kwargs)
        component.build_vector_store()

        # Switch to Retrieve mode
        component.set(mode="Retrieve", search_query="quick fox")
        results = component.search_documents()

        assert len(results) == 3
        # Results should be diverse but relevant
        assert any("fox" in result.text.lower() for result in results)
        mock_mmr_search.assert_called_once_with(query="quick fox", k=3)

        # Test with different settings
        component.set(number_of_results=2)
        mock_mmr_search.return_value = mock_docs[:2]
        diverse_results = component.search_documents()
        assert len(diverse_results) == 2

    @patch("langchain_chroma.Chroma.similarity_search")
    @patch("langchain_chroma.Chroma.max_marginal_relevance_search")
    def test_search_with_different_types(
        self,
        mock_mmr_search,
        mock_similarity_search,
        component_class: type[LocalDBComponent],
        default_kwargs: dict[str, Any],
    ) -> None:
        """Test search with different search types."""
        test_data = [
            "The quick brown fox jumps over the lazy dog",
            "Python is a popular programming language",
            "Machine learning models process data",
        ]
        default_kwargs["ingest_data"] = [Data(text=text) for text in test_data]
        default_kwargs["number_of_results"] = 2

        # Mock the search methods to return documents
        from langchain_core.documents import Document

        mock_similarity_docs = [
            Document(page_content="Python is a popular programming language"),
            Document(page_content="Machine learning models process data"),
        ]
        mock_similarity_search.return_value = mock_similarity_docs

        mock_mmr_docs = [
            Document(page_content="Python is a popular programming language"),
            Document(page_content="The quick brown fox jumps over the lazy dog"),
        ]
        mock_mmr_search.return_value = mock_mmr_docs

        component: LocalDBComponent = component_class().set(**default_kwargs)
        component.build_vector_store()

        # Switch to Retrieve mode and test similarity search
        component.set(mode="Retrieve", search_type="Similarity", search_query="programming languages")
        similarity_results = component.search_documents()
        assert len(similarity_results) == 2
        assert any("python" in result.text.lower() for result in similarity_results)
        mock_similarity_search.assert_called_once_with(query="programming languages", k=2)

        # Test MMR search
        component.set(search_type="MMR", search_query="programming languages")
        mmr_results = component.search_documents()
        assert len(mmr_results) == 2
        mock_mmr_search.assert_called_once_with(query="programming languages", k=2)

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

    @patch("langchain_chroma.Chroma.get")
    @patch("langchain_chroma.Chroma._collection")
    def test_duplicate_handling(
        self,
        mock_collection,
        mock_get,
        component_class: type[LocalDBComponent],
        default_kwargs: dict[str, Any],
    ) -> None:
        """Test handling of duplicate documents."""
        # Create test data with duplicates
        test_data = [
            Data(text_key="text", data={"text": "This is a test document"}),
            Data(text_key="text", data={"text": "This is a test document"}),  # Duplicate with exact same data
            Data(text_key="text", data={"text": "This is another document"}),
        ]
        default_kwargs["ingest_data"] = test_data
        default_kwargs["allow_duplicates"] = False
        default_kwargs["limit"] = 100  # Set a high enough limit to get all documents

        # Mock the get method to return documents
        mock_get.return_value = {
            "documents": ["This is a test document", "This is a test document", "This is another document"],
            "metadatas": [{}, {}, {}],
            "ids": ["1", "2", "3"],
        }

        # Mock collection count
        mock_collection.count.return_value = 3

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

        # Get all documents
        results = vector_store.get(limit=100)
        documents = results["documents"]

        # The documents are returned in a list structure
        assert len(documents) == 3  # All documents are added, even duplicates

        # Count unique texts
        unique_texts = set(documents)
        assert len(unique_texts) == 2  # Should have 2 unique texts

        # Test with allow_duplicates=True
        test_data = [
            Data(text_key="text", data={"text": "This is a test document"}),
            Data(text_key="text", data={"text": "This is a test document"}),  # Duplicate
        ]
        default_kwargs["ingest_data"] = test_data
        default_kwargs["allow_duplicates"] = True
        default_kwargs["collection_name"] = "test_collection_2"  # Use a different collection name

        # Mock for the second test
        mock_get.return_value = {
            "documents": ["This is a test document", "This is a test document"],
            "metadatas": [{}, {}],
            "ids": ["1", "2"],
        }
        mock_collection.count.return_value = 2

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

        # Get all documents
        results = vector_store.get(limit=100)
        documents = results["documents"]

        # With allow_duplicates=True, we should have both documents
        assert len(documents) == 2
        assert all("test document" in doc for doc in documents)

        # Verify that we have the expected number of documents
        assert vector_store._collection.count() == 2

    def test_build_config_update(self, component_class: type[LocalDBComponent]) -> None:
        """Test the update_build_config method."""
        component = component_class()

        # Test mode=Ingest
        build_config = {
            "ingest_data": {"show": False},
            "collection_name": {"show": False},
            "persist": {"show": False},
            "persist_directory": {"show": False},
            "embedding": {"show": False},
            "allow_duplicates": {"show": False},
            "limit": {"show": False},
            "search_query": {"show": False},
            "search_type": {"show": False},
            "number_of_results": {"show": False},
            "existing_collections": {"show": False},
        }

        updated_config = component.update_build_config(build_config, "Ingest", "mode")

        assert updated_config["ingest_data"]["show"] is True
        assert updated_config["collection_name"]["show"] is True
        assert updated_config["persist"]["show"] is True
        assert updated_config["search_query"]["show"] is False

        # Test mode=Retrieve
        updated_config = component.update_build_config(build_config, "Retrieve", "mode")

        assert updated_config["search_query"]["show"] is True
        assert updated_config["search_type"]["show"] is True
        assert updated_config["number_of_results"]["show"] is True
        assert updated_config["existing_collections"]["show"] is True
        assert updated_config["collection_name"]["show"] is False

        # Test persist=True/False
        build_config = {"persist_directory": {"show": False}}
        # Use keyword arguments to fix FBT003
        updated_config = component.update_build_config(build_config, field_value=True, field_name="persist")
        assert updated_config["persist_directory"]["show"] is True

        updated_config = component.update_build_config(build_config, field_value=False, field_name="persist")
        assert updated_config["persist_directory"]["show"] is False

        # Test existing_collections update
        # Fix the dict entry type issue
        build_config = {"collection_name": {"value": "old_name", "show": False}}
        updated_config = component.update_build_config(build_config, "new_collection", "existing_collections")
        assert updated_config["collection_name"]["value"] == "new_collection"

    @patch("lfx.components.vectorstores.local_db.LocalDBComponent.list_existing_collections")
    def test_list_existing_collections(self, mock_list: MagicMock, component_class: type[LocalDBComponent]) -> None:
        """Test the list_existing_collections method."""
        mock_list.return_value = ["collection1", "collection2", "collection3"]

        component = component_class()
        collections = component.list_existing_collections()

        assert collections == ["collection1", "collection2", "collection3"]
        mock_list.assert_called_once()

    def test_build_vector_store_disabled_in_astra_cloud(
        self, component_class: type[LocalDBComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test that build_vector_store raises an error when ASTRA_CLOUD_DISABLE_COMPONENT is true."""
        with patch.dict(os.environ, {"ASTRA_CLOUD_DISABLE_COMPONENT": "true"}):
            component: LocalDBComponent = component_class().set(**default_kwargs)

            with pytest.raises(ValueError) as exc_info:  # noqa: PT011
                component.build_vector_store()

            # Verify the error message mentions cloud/astra
            error_msg = str(exc_info.value).lower()
            assert "astra" in error_msg or "cloud" in error_msg or "s3" in error_msg

    def test_build_vector_store_works_when_not_in_astra_cloud(
        self, component_class: type[LocalDBComponent], default_kwargs: dict[str, Any]
    ) -> None:
        """Test that build_vector_store works normally when ASTRA_CLOUD_DISABLE_COMPONENT is false."""
        with patch.dict(os.environ, {"ASTRA_CLOUD_DISABLE_COMPONENT": "false"}):
            component: LocalDBComponent = component_class().set(**default_kwargs)

            # Should work without raising cloud-related errors
            vector_store = component.build_vector_store()
            assert vector_store is not None
