import importlib.util
import json
import uuid
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from lfx.base.knowledge_bases.knowledge_base_utils import get_knowledge_bases
from lfx.components.files_and_knowledge.retrieval import KnowledgeBaseComponent

from tests.base import ComponentTestBaseWithClient

# langchain-ibm is gated to python_version<'3.14' in pyproject.toml because
# upstream pins exclude 3.14. Tests that patch langchain_ibm.* directly must
# be skipped when the package is unavailable.
_HAS_LANGCHAIN_IBM = importlib.util.find_spec("langchain_ibm") is not None
_NO_LANGCHAIN_IBM_REASON = "langchain-ibm not available (gated to python_version<'3.14')"


class TestKnowledgeBaseComponent(ComponentTestBaseWithClient):
    @pytest.fixture
    def component_class(self):
        """Return the component class to test."""
        return KnowledgeBaseComponent

    @pytest.fixture(autouse=True)
    def mock_knowledge_base_path(self, tmp_path):
        """Mock the knowledge base root path directly."""
        with patch("langflow.components.knowledge_bases.retrieval._KNOWLEDGE_BASES_ROOT_PATH", tmp_path):
            yield

    @pytest.fixture
    def default_kwargs(self, tmp_path, active_user):
        """Return default kwargs for component instantiation."""
        # Create knowledge base directory structure
        kb_name = "test_kb"
        kb_path = tmp_path / active_user.username / kb_name
        kb_path.mkdir(parents=True, exist_ok=True)

        # Create embedding metadata file
        metadata = {
            "embedding_provider": "HuggingFace",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "api_key": None,
            "api_key_used": False,
            "chunk_size": 1000,
            "created_at": "2024-01-01T00:00:00Z",
        }
        (kb_path / "embedding_metadata.json").write_text(json.dumps(metadata))

        return {
            "knowledge_base": kb_name,
            "kb_root_path": str(tmp_path),
            "api_key": None,
            "search_query": "",
            "top_k": 5,
            "include_embeddings": True,
            "_user_id": active_user.id,
        }

    @pytest.fixture
    def file_names_mapping(self):
        """Return file names mapping for version testing."""
        # This is a new component, so it doesn't exist in older versions
        return []

    async def test_get_knowledge_bases(self, tmp_path, active_user):
        """Test getting list of knowledge bases."""
        # Create additional test directories
        (tmp_path / active_user.username / "kb1").mkdir(parents=True, exist_ok=True)
        (tmp_path / active_user.username / "kb2").mkdir(parents=True, exist_ok=True)
        (tmp_path / active_user.username / ".hidden").mkdir(parents=True, exist_ok=True)  # Should be ignored

        kb_list = await get_knowledge_bases(tmp_path, user_id=active_user.id)

        assert "test_kb" in kb_list
        assert "kb1" in kb_list
        assert "kb2" in kb_list
        assert ".hidden" not in kb_list

    async def test_update_build_config(self, component_class, default_kwargs, tmp_path, active_user):
        """Test updating build configuration."""
        component = component_class(**default_kwargs)

        # Create additional KB directories
        (tmp_path / active_user.username / "kb1").mkdir(parents=True, exist_ok=True)
        (tmp_path / active_user.username / "kb2").mkdir(parents=True, exist_ok=True)

        build_config = {"knowledge_base": {"value": "test_kb", "options": []}}

        result = await component.update_build_config(build_config, None, "knowledge_base")

        assert "test_kb" in result["knowledge_base"]["options"]
        assert "kb1" in result["knowledge_base"]["options"]
        assert "kb2" in result["knowledge_base"]["options"]

    async def test_update_build_config_invalid_kb(self, component_class, default_kwargs):
        """Test updating build config when selected KB is not available."""
        component = component_class(**default_kwargs)

        build_config = {"knowledge_base": {"value": "nonexistent_kb", "options": ["test_kb"]}}

        result = await component.update_build_config(build_config, None, "knowledge_base")

        assert result["knowledge_base"]["value"] is None

    def test_get_kb_metadata_success(self, component_class, default_kwargs, active_user):
        """Test successful metadata loading."""
        component = component_class(**default_kwargs)
        kb_path = Path(default_kwargs["kb_root_path"]) / active_user.username / default_kwargs["knowledge_base"]

        with patch("langflow.components.knowledge_bases.retrieval.decrypt_api_key") as mock_decrypt:
            mock_decrypt.return_value = "decrypted_key"

            metadata = component._get_kb_metadata(kb_path)

        assert metadata["embedding_provider"] == "HuggingFace"
        assert metadata["embedding_model"] == "sentence-transformers/all-MiniLM-L6-v2"
        assert "chunk_size" in metadata

    def test_get_kb_metadata_no_file(self, component_class, default_kwargs, tmp_path, active_user):
        """Test metadata loading when file doesn't exist."""
        component = component_class(**default_kwargs)
        nonexistent_path = tmp_path / active_user.username / "nonexistent"
        nonexistent_path.mkdir(parents=True, exist_ok=True)

        metadata = component._get_kb_metadata(nonexistent_path)

        assert metadata == {}

    def test_get_kb_metadata_json_error(self, component_class, default_kwargs, tmp_path, active_user):
        """Test metadata loading with invalid JSON."""
        component = component_class(**default_kwargs)
        kb_path = tmp_path / active_user.username / "invalid_json_kb"
        kb_path.mkdir(parents=True, exist_ok=True)

        # Create invalid JSON file
        (kb_path / "embedding_metadata.json").write_text("invalid json content")

        metadata = component._get_kb_metadata(kb_path)

        assert metadata == {}

    def test_get_kb_metadata_decrypt_error(self, component_class, default_kwargs, tmp_path, active_user):
        """Test metadata loading with decryption error."""
        component = component_class(**default_kwargs)
        kb_path = tmp_path / active_user.username / "decrypt_error_kb"
        kb_path.mkdir(parents=True, exist_ok=True)

        # Create metadata with encrypted key
        metadata = {
            "embedding_provider": "OpenAI",
            "embedding_model": "text-embedding-ada-002",
            "api_key": "encrypted_key",  # pragma:allowlist secret
            "chunk_size": 1000,
        }
        (kb_path / "embedding_metadata.json").write_text(json.dumps(metadata))

        with patch("langflow.components.knowledge_bases.retrieval.decrypt_api_key") as mock_decrypt:
            mock_decrypt.side_effect = ValueError("Decryption failed")

            result = component._get_kb_metadata(kb_path)

        assert result["api_key"] is None

    @patch("langchain_huggingface.HuggingFaceEmbeddings")
    def test_build_embeddings_huggingface(self, mock_hf_embeddings, component_class, default_kwargs):
        """Test building HuggingFace embeddings."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "HuggingFace",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_hf_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(metadata)

        mock_hf_embeddings.assert_called_once_with(model="sentence-transformers/all-MiniLM-L6-v2")
        assert result == mock_embeddings

    @patch("langchain_openai.OpenAIEmbeddings")
    def test_build_embeddings_openai(self, mock_openai_embeddings, component_class, default_kwargs):
        """Test building OpenAI embeddings."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "OpenAI",
            "embedding_model": "text-embedding-ada-002",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_openai_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(metadata, api_key="test-api-key")

        mock_openai_embeddings.assert_called_once_with(
            model="text-embedding-ada-002",
            api_key="test-api-key",  # pragma:allowlist secret
            chunk_size=1000,  # pragma:allowlist secret
        )
        assert result == mock_embeddings

    def test_build_embeddings_openai_no_key(self, component_class, default_kwargs):
        """Test building OpenAI embeddings without API key raises error."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "OpenAI",
            "embedding_model": "text-embedding-ada-002",
            "chunk_size": 1000,
        }

        with pytest.raises(ValueError, match="OpenAI API key is required"):
            component._build_embeddings(metadata)

    @patch("langchain_cohere.CohereEmbeddings")
    def test_build_embeddings_cohere(self, mock_cohere_embeddings, component_class, default_kwargs):
        """Test building Cohere embeddings."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Cohere",
            "embedding_model": "embed-english-v3.0",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_cohere_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(metadata, api_key="test-api-key")

        mock_cohere_embeddings.assert_called_once_with(
            model="embed-english-v3.0",
            cohere_api_key="test-api-key",  # pragma:allowlist secret
        )  # pragma:allowlist secret
        assert result == mock_embeddings

    def test_build_embeddings_cohere_no_key(self, component_class, default_kwargs):
        """Test building Cohere embeddings without API key raises error."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Cohere",
            "embedding_model": "embed-english-v3.0",
            "chunk_size": 1000,
        }

        with pytest.raises(ValueError, match="Cohere API key is required"):
            component._build_embeddings(metadata)

    def test_build_embeddings_custom_not_supported(self, component_class, default_kwargs):
        """Test building custom embeddings raises NotImplementedError."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Custom",
            "embedding_model": "custom-model",
        }

        with pytest.raises(NotImplementedError, match="Custom embedding models not yet supported"):
            component._build_embeddings(metadata)

    def test_build_embeddings_unsupported_provider(self, component_class, default_kwargs):
        """Test building embeddings with unsupported provider raises NotImplementedError."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "UnsupportedProvider",
            "embedding_model": "some-model",
            "api_key": "test-key",  # pragma:allowlist secret
        }  # pragma:allowlist secret

        with pytest.raises(NotImplementedError, match="Embedding provider 'UnsupportedProvider' is not supported"):
            component._build_embeddings(metadata)

    @patch("langchain_google_genai.GoogleGenerativeAIEmbeddings")
    def test_build_embeddings_google_generative_ai(self, mock_google_embeddings, component_class, default_kwargs):
        """Test building Google Generative AI embeddings."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Google Generative AI",
            "embedding_model": "models/embedding-001",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_google_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(metadata, api_key="test-google-key")

        mock_google_embeddings.assert_called_once_with(
            model="models/embedding-001",
            google_api_key="test-google-key",  # pragma:allowlist secret
        )
        assert result == mock_embeddings

    def test_build_embeddings_google_no_key(self, component_class, default_kwargs):
        """Test building Google Generative AI embeddings without API key raises error."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Google Generative AI",
            "embedding_model": "models/embedding-001",
            "chunk_size": 1000,
        }

        with pytest.raises(ValueError, match="Google API key is required"):
            component._build_embeddings(metadata)

    @patch("langchain_ollama.OllamaEmbeddings")
    def test_build_embeddings_ollama(self, mock_ollama_embeddings, component_class, default_kwargs):
        """Test building Ollama embeddings."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Ollama",
            "embedding_model": "nomic-embed-text",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_ollama_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(
            metadata,
            provider_vars={"OLLAMA_BASE_URL": "http://localhost:11434"},
        )

        mock_ollama_embeddings.assert_called_once_with(
            model="nomic-embed-text",
            base_url="http://localhost:11434",
        )
        assert result == mock_embeddings

    @pytest.mark.skipif(not _HAS_LANGCHAIN_IBM, reason=_NO_LANGCHAIN_IBM_REASON)
    @patch("langchain_ibm.WatsonxEmbeddings")
    def test_build_embeddings_watsonx(self, mock_watsonx_embeddings, component_class, default_kwargs):
        """Test building IBM WatsonX embeddings."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "IBM WatsonX",
            "embedding_model": "ibm/slate-125m-english-rtrvr-v2",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_watsonx_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(
            metadata,
            api_key="test-watsonx-key",
            provider_vars={
                "WATSONX_APIKEY": "test-watsonx-key",  # pragma:allowlist secret
                "WATSONX_PROJECT_ID": "test-project-id",
                "WATSONX_URL": "https://us-south.ml.cloud.ibm.com",
            },
        )

        mock_watsonx_embeddings.assert_called_once_with(
            model_id="ibm/slate-125m-english-rtrvr-v2",
            apikey="test-watsonx-key",  # pragma:allowlist secret
            project_id="test-project-id",
            url="https://us-south.ml.cloud.ibm.com",
        )
        assert result == mock_embeddings

    @pytest.mark.skipif(not _HAS_LANGCHAIN_IBM, reason=_NO_LANGCHAIN_IBM_REASON)
    def test_build_embeddings_watsonx_no_key(self, component_class, default_kwargs):
        """Test building IBM WatsonX embeddings without API key raises error."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "IBM WatsonX",
            "embedding_model": "ibm/slate-125m-english-rtrvr-v2",
            "chunk_size": 1000,
        }

        with pytest.raises(ValueError, match="IBM WatsonX API key is required"):
            component._build_embeddings(metadata)

    @patch("langchain_openai.OpenAIEmbeddings")
    async def test_resolve_api_key_global_fallback(self, mock_openai_embeddings, component_class, default_kwargs):
        """Test that retrieve_data resolves the global API key for OpenAI."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "OpenAI",
            "embedding_model": "text-embedding-ada-002",
            "chunk_size": 1000,
        }

        mock_embeddings = MagicMock()
        mock_openai_embeddings.return_value = mock_embeddings

        # The async _resolve_api_key should find the global key
        with patch.object(component, "_resolve_api_key", return_value="global-openai-key"):
            result = component._build_embeddings(metadata, api_key="global-openai-key")

        mock_openai_embeddings.assert_called_once_with(
            model="text-embedding-ada-002",
            api_key="global-openai-key",  # pragma:allowlist secret
            chunk_size=1000,
        )
        assert result == mock_embeddings

    def test_build_embeddings_with_explicit_api_key(self, component_class, default_kwargs):
        """Test that an explicit API key is used when passed."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "OpenAI",
            "embedding_model": "text-embedding-ada-002",
            "chunk_size": 1000,
        }

        with patch("langchain_openai.OpenAIEmbeddings") as mock_openai:
            mock_embeddings = MagicMock()
            mock_openai.return_value = mock_embeddings

            component._build_embeddings(metadata, api_key="user-provided-key")

            mock_openai.assert_called_once_with(
                model="text-embedding-ada-002",
                api_key="user-provided-key",  # pragma:allowlist secret
                chunk_size=1000,
            )

    async def test_retrieve_data_no_metadata(self, component_class, default_kwargs, tmp_path, active_user):
        """Test retrieving data when metadata is missing."""
        # Remove metadata file
        kb_path = tmp_path / active_user.username / default_kwargs["knowledge_base"]
        metadata_file = kb_path / "embedding_metadata.json"
        if metadata_file.exists():
            metadata_file.unlink()

        component = component_class(**default_kwargs)

        with pytest.raises(ValueError, match="Metadata not found for knowledge base"):
            await component.retrieve_data()

    def test_retrieve_data_path_construction(self, component_class, default_kwargs):
        """Test that retrieve_data constructs the correct paths."""
        component = component_class(**default_kwargs)

        # Test that the component correctly builds the KB path

        assert component.kb_root_path == default_kwargs["kb_root_path"]
        assert component.knowledge_base == default_kwargs["knowledge_base"]

        # Test that paths are correctly expanded
        expanded_path = Path(component.kb_root_path).expanduser()
        assert expanded_path.exists()  # tmp_path should exist

        # Verify method exists with correct parameters
        assert hasattr(component, "retrieve_data")
        assert hasattr(component, "search_query")
        assert hasattr(component, "top_k")
        assert hasattr(component, "include_embeddings")

    async def test_retrieve_data_method_exists(self, component_class, default_kwargs, active_user):
        """Test that retrieve_data method exists and can be called."""
        component = component_class(**default_kwargs)

        # Just verify the method exists and has the right signature
        assert hasattr(component, "retrieve_data"), "Component should have retrieve_data method"

        # Build a mock Chroma that returns results in the expected format
        mock_doc = MagicMock()
        mock_doc.page_content = "test content"
        mock_doc.metadata = {"_id": "doc1", "source": "test"}

        mock_chroma_instance = MagicMock()
        mock_chroma_instance.similarity_search.return_value = [mock_doc]

        mock_user = MagicMock()
        mock_user.username = active_user.username

        # Mock all external calls to avoid integration issues
        with (
            patch.object(component, "_get_kb_metadata") as mock_get_metadata,
            patch.object(component, "_build_embeddings") as mock_build_embeddings,
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch("lfx.components.files_and_knowledge.retrieval.get_user_by_id", return_value=mock_user),
            patch(
                "lfx.components.files_and_knowledge.retrieval._get_knowledge_bases_root_path",
                return_value=Path(default_kwargs["kb_root_path"]),
            ),
            patch("chromadb.api.client.SharedSystemClient.clear_system_cache"),
            patch("lfx.components.files_and_knowledge.retrieval.Chroma", return_value=mock_chroma_instance),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)
            mock_get_metadata.return_value = {"embedding_provider": "HuggingFace", "embedding_model": "test-model"}
            mock_build_embeddings.return_value = MagicMock()

            result = await component.retrieve_data()

            # Verify internal methods were called
            mock_get_metadata.assert_called_once()
            mock_build_embeddings.assert_called_once()
            assert len(result) == 1

    def test_include_embeddings_parameter(self, component_class, default_kwargs):
        """Test that include_embeddings parameter is properly set."""
        # Test with embeddings enabled
        default_kwargs["include_embeddings"] = True
        component = component_class(**default_kwargs)
        assert component.include_embeddings is True

        # Test with embeddings disabled
        default_kwargs["include_embeddings"] = False
        component = component_class(**default_kwargs)
        assert component.include_embeddings is False

    # --- _resolve_provider_variables tests ---

    async def test_resolve_provider_variables_empty_provider_vars(self, component_class, default_kwargs):
        """Test _resolve_provider_variables when provider has no variables defined."""
        component = component_class(**default_kwargs)

        with patch(
            "lfx.components.files_and_knowledge.retrieval.get_provider_all_variables",
            return_value=[],
        ):
            result = await component._resolve_provider_variables("Ollama")

        assert result == {}

    async def test_resolve_provider_variables_variable_service_returns_none(self, component_class, default_kwargs):
        """Test _resolve_provider_variables when variable_service is None."""
        component = component_class(**default_kwargs)

        with (
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_provider_all_variables",
                return_value=[{"variable_key": "OLLAMA_BASE_URL"}],
            ),
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_variable_service",
                return_value=None,
            ),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            result = await component._resolve_provider_variables("Ollama")

        assert result == {}

    async def test_resolve_provider_variables_user_id_is_none(self, component_class, default_kwargs):
        """Test _resolve_provider_variables when user_id is None."""
        default_kwargs["_user_id"] = None
        component = component_class(**default_kwargs)

        # Set a mock vertex so user_id property returns None
        # instead of the string "None" from PlaceholderGraph
        mock_vertex = MagicMock()
        mock_vertex.graph.user_id = None
        component._vertex = mock_vertex

        with patch(
            "lfx.components.files_and_knowledge.retrieval.get_provider_all_variables",
            return_value=[{"variable_key": "OLLAMA_BASE_URL"}],
        ):
            result = await component._resolve_provider_variables("Ollama")

        assert result == {}

    async def test_resolve_provider_variables_user_id_as_string(self, component_class, default_kwargs):
        """Test _resolve_provider_variables when user_id is a string UUID."""
        user_uuid = uuid.uuid4()
        default_kwargs["_user_id"] = str(user_uuid)
        component = component_class(**default_kwargs)

        mock_variable_service = AsyncMock()
        mock_variable_service.get_variable = AsyncMock(return_value="http://localhost:11434")

        with (
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_provider_all_variables",
                return_value=[{"variable_key": "OLLAMA_BASE_URL"}],
            ),
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_variable_service",
                return_value=mock_variable_service,
            ),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            result = await component._resolve_provider_variables("Ollama")

        assert result == {"OLLAMA_BASE_URL": "http://localhost:11434"}
        # Verify the user_id was correctly converted to UUID
        call_kwargs = mock_variable_service.get_variable.call_args[1]
        assert call_kwargs["user_id"] == user_uuid

    async def test_resolve_provider_variables_lookup_falls_back_to_env(
        self, component_class, default_kwargs, monkeypatch
    ):
        """Test _resolve_provider_variables falls back to env var on service error."""
        component = component_class(**default_kwargs)

        mock_variable_service = AsyncMock()
        mock_variable_service.get_variable = AsyncMock(side_effect=ValueError("Not found"))
        monkeypatch.setenv("OLLAMA_BASE_URL", "http://env-fallback:11434")

        with (
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_provider_all_variables",
                return_value=[{"variable_key": "OLLAMA_BASE_URL"}],
            ),
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_variable_service",
                return_value=mock_variable_service,
            ),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            result = await component._resolve_provider_variables("Ollama")

        assert result == {"OLLAMA_BASE_URL": "http://env-fallback:11434"}

    # --- _resolve_api_key tests ---

    async def test_resolve_api_key_unknown_provider(self, component_class, default_kwargs):
        """Test _resolve_api_key when provider is not in provider_variable_map."""
        component = component_class(**default_kwargs)

        with patch(
            "lfx.components.files_and_knowledge.retrieval.get_model_provider_variable_mapping",
            return_value={},
        ):
            result = await component._resolve_api_key("UnknownProvider")

        assert result is None

    async def test_resolve_api_key_user_id_is_none(self, component_class, default_kwargs):
        """Test _resolve_api_key when user_id is None."""
        default_kwargs["_user_id"] = None
        component = component_class(**default_kwargs)

        # Set a mock vertex so user_id property returns None
        # instead of the string "None" from PlaceholderGraph
        mock_vertex = MagicMock()
        mock_vertex.graph.user_id = None
        component._vertex = mock_vertex

        with patch(
            "lfx.components.files_and_knowledge.retrieval.get_model_provider_variable_mapping",
            return_value={"OpenAI": "OPENAI_API_KEY"},
        ):
            result = await component._resolve_api_key("OpenAI")

        assert result is None

    async def test_resolve_api_key_variable_service_is_none(self, component_class, default_kwargs):
        """Test _resolve_api_key when variable_service returns None."""
        component = component_class(**default_kwargs)

        with (
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_model_provider_variable_mapping",
                return_value={"OpenAI": "OPENAI_API_KEY"},
            ),
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_variable_service",
                return_value=None,
            ),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            result = await component._resolve_api_key("OpenAI")

        assert result is None

    async def test_resolve_api_key_variable_service_raises(self, component_class, default_kwargs):
        """Test _resolve_api_key returns None when variable_service.get_variable raises."""
        component = component_class(**default_kwargs)

        mock_variable_service = AsyncMock()
        mock_variable_service.get_variable = AsyncMock(side_effect=ValueError("Not found"))

        with (
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_model_provider_variable_mapping",
                return_value={"OpenAI": "OPENAI_API_KEY"},
            ),
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_variable_service",
                return_value=mock_variable_service,
            ),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            result = await component._resolve_api_key("OpenAI")

        assert result is None

    # --- _build_embeddings edge case tests ---

    @patch("langchain_ollama.OllamaEmbeddings")
    def test_build_embeddings_ollama_without_base_url(self, mock_ollama_embeddings, component_class, default_kwargs):
        """Test building Ollama embeddings without base_url (empty provider_vars)."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Ollama",
            "embedding_model": "nomic-embed-text",
        }

        mock_embeddings = MagicMock()
        mock_ollama_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(metadata, provider_vars={})

        # Should be called without base_url kwarg
        mock_ollama_embeddings.assert_called_once_with(model="nomic-embed-text")
        assert result == mock_embeddings

    @patch("langchain_ollama.OllamaEmbeddings")
    def test_build_embeddings_ollama_no_provider_vars(self, mock_ollama_embeddings, component_class, default_kwargs):
        """Test building Ollama embeddings with provider_vars=None."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "Ollama",
            "embedding_model": "nomic-embed-text",
        }

        mock_embeddings = MagicMock()
        mock_ollama_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(metadata, provider_vars=None)

        mock_ollama_embeddings.assert_called_once_with(model="nomic-embed-text")
        assert result == mock_embeddings

    @pytest.mark.skipif(not _HAS_LANGCHAIN_IBM, reason=_NO_LANGCHAIN_IBM_REASON)
    @patch("langchain_ibm.WatsonxEmbeddings")
    def test_build_embeddings_watsonx_api_key_from_provider_vars(
        self, mock_watsonx_embeddings, component_class, default_kwargs
    ):
        """Test WatsonX uses api_key from provider_vars fallback when api_key param is None."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "IBM WatsonX",
            "embedding_model": "ibm/slate-125m-english-rtrvr-v2",
        }

        mock_embeddings = MagicMock()
        mock_watsonx_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(
            metadata,
            api_key=None,
            provider_vars={
                "WATSONX_APIKEY": "vars-watsonx-key",  # pragma:allowlist secret
                "WATSONX_PROJECT_ID": "project-123",
                "WATSONX_URL": "https://us-south.ml.cloud.ibm.com",
            },
        )

        mock_watsonx_embeddings.assert_called_once_with(
            model_id="ibm/slate-125m-english-rtrvr-v2",
            apikey="vars-watsonx-key",  # pragma:allowlist secret
            project_id="project-123",
            url="https://us-south.ml.cloud.ibm.com",
        )
        assert result == mock_embeddings

    @pytest.mark.skipif(not _HAS_LANGCHAIN_IBM, reason=_NO_LANGCHAIN_IBM_REASON)
    @patch("langchain_ibm.WatsonxEmbeddings")
    def test_build_embeddings_watsonx_partial_vars(self, mock_watsonx_embeddings, component_class, default_kwargs):
        """Test WatsonX with only apikey, no project_id or url."""
        component = component_class(**default_kwargs)

        metadata = {
            "embedding_provider": "IBM WatsonX",
            "embedding_model": "ibm/slate-125m-english-rtrvr-v2",
        }

        mock_embeddings = MagicMock()
        mock_watsonx_embeddings.return_value = mock_embeddings

        result = component._build_embeddings(
            metadata,
            api_key="only-api-key",
            provider_vars={},
        )

        # project_id and url should be omitted from kwargs
        mock_watsonx_embeddings.assert_called_once_with(
            model_id="ibm/slate-125m-english-rtrvr-v2",
            apikey="only-api-key",  # pragma:allowlist secret
        )
        assert result == mock_embeddings

    def test_build_embeddings_empty_metadata(self, component_class, default_kwargs):
        """Test _build_embeddings with empty metadata dict (no provider, no model)."""
        component = component_class(**default_kwargs)

        with pytest.raises(NotImplementedError, match="Embedding provider 'None' is not supported"):
            component._build_embeddings({})

    # --- retrieve_data integration edge case tests ---

    async def test_retrieve_data_user_id_is_none(self, component_class, default_kwargs):
        """Test retrieve_data raises when user_id is None."""
        default_kwargs["_user_id"] = None
        component = component_class(**default_kwargs)

        # Set a mock vertex so user_id property returns None
        # instead of the string "None" from PlaceholderGraph
        mock_vertex = MagicMock()
        mock_vertex.graph.user_id = None
        component._vertex = mock_vertex

        with (
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            with pytest.raises(ValueError, match="User ID is required"):
                await component.retrieve_data()

    async def test_retrieve_data_user_not_found(self, component_class, default_kwargs):
        """Test retrieve_data raises when user is not found in DB."""
        component = component_class(**default_kwargs)

        with (
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch(
                "lfx.components.files_and_knowledge.retrieval.get_user_by_id",
                return_value=None,
            ),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)

            with pytest.raises(ValueError, match=r"User with ID .* not found"):
                await component.retrieve_data()

    async def test_retrieve_data_with_search_query(self, component_class, default_kwargs, active_user):
        """Test retrieve_data with a populated search_query uses similarity_search_with_score."""
        default_kwargs["search_query"] = "find me something"
        component = component_class(**default_kwargs)

        mock_doc = MagicMock()
        mock_doc.page_content = "matched content"
        mock_doc.metadata = {"_id": "doc1", "source": "test"}

        mock_chroma_instance = MagicMock()
        # similarity_search_with_score returns (doc, score) tuples
        mock_chroma_instance.similarity_search_with_score.return_value = [(mock_doc, 0.85)]

        mock_user = MagicMock()
        mock_user.username = active_user.username

        with (
            patch.object(component, "_get_kb_metadata") as mock_get_metadata,
            patch.object(component, "_build_embeddings") as mock_build_embeddings,
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch("lfx.components.files_and_knowledge.retrieval.get_user_by_id", return_value=mock_user),
            patch(
                "lfx.components.files_and_knowledge.retrieval._get_knowledge_bases_root_path",
                return_value=Path(default_kwargs["kb_root_path"]),
            ),
            patch("chromadb.api.client.SharedSystemClient.clear_system_cache"),
            patch("lfx.components.files_and_knowledge.retrieval.Chroma", return_value=mock_chroma_instance),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)
            mock_get_metadata.return_value = {"embedding_provider": "HuggingFace", "embedding_model": "test-model"}
            mock_build_embeddings.return_value = MagicMock()

            result = await component.retrieve_data()

            # Verify similarity_search_with_score was used (not similarity_search)
            mock_chroma_instance.similarity_search_with_score.assert_called_once_with(query="find me something", k=5)
            mock_chroma_instance.similarity_search.assert_not_called()
            assert len(result) == 1

    async def test_retrieve_data_without_search_query(self, component_class, default_kwargs, active_user):
        """Test retrieve_data with empty search_query uses similarity_search."""
        default_kwargs["search_query"] = ""
        component = component_class(**default_kwargs)

        mock_doc = MagicMock()
        mock_doc.page_content = "all content"
        mock_doc.metadata = {"_id": "doc1", "source": "test"}

        mock_chroma_instance = MagicMock()
        mock_chroma_instance.similarity_search.return_value = [mock_doc]

        mock_user = MagicMock()
        mock_user.username = active_user.username

        with (
            patch.object(component, "_get_kb_metadata") as mock_get_metadata,
            patch.object(component, "_build_embeddings") as mock_build_embeddings,
            patch("lfx.components.files_and_knowledge.retrieval.session_scope") as mock_session_scope,
            patch("lfx.components.files_and_knowledge.retrieval.get_user_by_id", return_value=mock_user),
            patch(
                "lfx.components.files_and_knowledge.retrieval._get_knowledge_bases_root_path",
                return_value=Path(default_kwargs["kb_root_path"]),
            ),
            patch("chromadb.api.client.SharedSystemClient.clear_system_cache"),
            patch("lfx.components.files_and_knowledge.retrieval.Chroma", return_value=mock_chroma_instance),
        ):
            mock_session = AsyncMock()
            mock_session_scope.return_value.__aenter__ = AsyncMock(return_value=mock_session)
            mock_session_scope.return_value.__aexit__ = AsyncMock(return_value=False)
            mock_get_metadata.return_value = {"embedding_provider": "HuggingFace", "embedding_model": "test-model"}
            mock_build_embeddings.return_value = MagicMock()

            result = await component.retrieve_data()

            # Verify similarity_search was used (not similarity_search_with_score)
            mock_chroma_instance.similarity_search.assert_called_once_with(query="", k=5)
            mock_chroma_instance.similarity_search_with_score.assert_not_called()
            assert len(result) == 1
