"""
Test Bedrock Vector Store helper functions and transformation.
"""

import pytest
from unittest.mock import Mock
import httpx

from tests.vector_store_tests.base_vector_store_test import BaseVectorStoreTest
from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj


class TestBedrockVectorStore(BaseVectorStoreTest):
    """
    Test the Bedrock vector store transformation functionality.
    """

    def get_base_create_vector_store_args(self) -> dict:
        """Must return the base create vector store args"""
        return {}

    def get_base_request_args(self):
        return {
            "vector_store_id": "T37J8R4WTM",
            "custom_llm_provider": "bedrock",
            "query": "what happens after we add a model",
        }

    def test_get_file_id_from_metadata(self):
        """Test that file_id is correctly extracted from metadata."""
        config = BedrockVectorStoreConfig()

        # Test with source URI
        metadata_with_uri = {
            "x-amz-bedrock-kb-source-uri": "https://www.litellm.ai",
            "x-amz-bedrock-kb-chunk-id": "1%3A0%3AjNYPg5YByRuP5PdK96co",
        }
        file_id = config._get_file_id_from_metadata(metadata_with_uri)
        assert file_id == "https://www.litellm.ai"

        # Test without source URI but with chunk ID
        metadata_without_uri = {
            "x-amz-bedrock-kb-chunk-id": "1%3A0%3AjNYPg5YByRuP5PdK96co"
        }
        file_id = config._get_file_id_from_metadata(metadata_without_uri)
        assert file_id == "bedrock-kb-1%3A0%3AjNYPg5YByRuP5PdK96co"

        # Test with empty metadata
        file_id = config._get_file_id_from_metadata({})
        assert file_id == "bedrock-kb-unknown"

    def test_get_filename_from_metadata(self):
        """Test that filename is correctly extracted from metadata."""
        config = BedrockVectorStoreConfig()

        # Test with source URI containing path
        metadata_with_path = {
            "x-amz-bedrock-kb-source-uri": "https://docs.litellm.ai/tutorial/setup.html"
        }
        filename = config._get_filename_from_metadata(metadata_with_path)
        assert filename == "setup.html"

        # Test with source URI without path (domain only)
        metadata_domain_only = {"x-amz-bedrock-kb-source-uri": "https://www.litellm.ai"}
        filename = config._get_filename_from_metadata(metadata_domain_only)
        assert filename == "www.litellm.ai"

        # Test without source URI but with data source ID
        metadata_without_uri = {"x-amz-bedrock-kb-data-source-id": "CCEJIRXXFI"}
        filename = config._get_filename_from_metadata(metadata_without_uri)
        assert filename == "bedrock-kb-document-CCEJIRXXFI"

        # Test with empty metadata
        filename = config._get_filename_from_metadata({})
        assert filename == "bedrock-kb-document-unknown"

    def test_get_attributes_from_metadata(self):
        """Test that attributes are correctly extracted from metadata."""
        config = BedrockVectorStoreConfig()

        # Test with full metadata
        metadata = {
            "x-amz-bedrock-kb-source-uri": "https://www.litellm.ai",
            "x-amz-bedrock-kb-chunk-id": "1%3A0%3AjNYPg5YByRuP5PdK96co",
            "x-amz-bedrock-kb-data-source-id": "CCEJIRXXFI",
        }
        attributes = config._get_attributes_from_metadata(metadata)
        assert attributes == metadata
        assert attributes is not metadata  # Should be a copy

        # Test with empty metadata
        attributes = config._get_attributes_from_metadata({})
        assert attributes == {}

        # Test with None
        attributes = config._get_attributes_from_metadata(None)
        assert attributes == {}


@pytest.mark.asyncio
async def test_bedrock_search_with_router():
    from litellm.router import Router

    # init router
    _router = Router(model_list=[])
    search_response = await _router.avector_store_search(
        query="what happens after we add a model",
        vector_store_id="T37J8R4WTM",
        custom_llm_provider="bedrock",
    )
    print(search_response)


@pytest.mark.asyncio
async def test_bedrock_search_with_credentials_managed_registry():
    """
    Test that the vector store search uses the credential accessor from the registry
    when AWS environment variables are not set, ensuring credentials are managed properly.
    """
    from unittest.mock import patch, MagicMock
    from litellm.router import Router
    from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
    from litellm.types.utils import CredentialItem
    from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
    from datetime import datetime, timezone
    import litellm

    # Store original registry and credential list
    original_registry = getattr(litellm, "vector_store_registry", None)
    original_credential_list = getattr(litellm, "credential_list", [])

    try:
        # Set up test AWS credentials in the credential system
        test_credentials = CredentialItem(
            credential_name="bedrock-litellm-website-knowledgebase",
            credential_info={
                "provider": "aws",
                "description": "Test AWS credentials for bedrock",
            },
            credential_values={
                "aws_access_key_id": "test_access_key",
                "aws_secret_access_key": "test_secret_key",
                "aws_region_name": "us-east-1",
            },
        )

        # Set up the credential list
        litellm.credential_list = [test_credentials]

        # Create vector store with credential reference
        vector_store = LiteLLM_ManagedVectorStore(
            vector_store_id="T37J8R4WTM",
            custom_llm_provider="bedrock",
            created_at=datetime.now(timezone.utc),
            updated_at=datetime.now(timezone.utc),
            litellm_credential_name="bedrock-litellm-website-knowledgebase",
        )

        # Set up registry
        registry = VectorStoreRegistry([vector_store])
        litellm.vector_store_registry = registry

        # Verify credentials can be retrieved from registry
        retrieved_credentials = registry.get_credentials_for_vector_store("T37J8R4WTM")
        assert retrieved_credentials, "Should retrieve credentials from registry"
        assert retrieved_credentials.get("aws_access_key_id") == "test_access_key"
        assert retrieved_credentials.get("aws_secret_access_key") == "test_secret_key"
        assert retrieved_credentials.get("aws_region_name") == "us-east-1"

        # Create router and perform search
        _router = Router(model_list=[])

        # Mock the credential injection process to verify it's called
        with patch.object(
            registry,
            "get_credentials_for_vector_store",
            wraps=registry.get_credentials_for_vector_store,
        ) as mock_get_creds:
            # Mock the actual search call to avoid making real API calls
            with patch(
                "litellm.vector_stores.main.base_llm_http_handler.vector_store_search_handler"
            ) as mock_handler:
                mock_handler.return_value = {
                    "data": [
                        {
                            "id": "test_result",
                            "text": "Mock search result",
                            "score": 0.9,
                            "metadata": {},
                        }
                    ]
                }

                search_response = await _router.avector_store_search(
                    query="what happens after we add a model",
                    vector_store_id="T37J8R4WTM",
                    custom_llm_provider="bedrock",
                )

                # Verify the search was called
                mock_handler.assert_called_once()
                call_kwargs = mock_handler.call_args[1]

                # Verify that the credential accessor was called with the correct vector store ID
                mock_get_creds.assert_called_with("T37J8R4WTM")

                # Verify the credentials were injected into the search call
                litellm_params = call_kwargs.get("litellm_params", {})

                # The key test: verify that credentials from the registry were used
                # Since we have a registry with credentials, they should be present in the params
                assert hasattr(
                    litellm_params, "aws_access_key_id"
                ), "aws_access_key_id should be in litellm_params"
                assert hasattr(
                    litellm_params, "aws_secret_access_key"
                ), "aws_secret_access_key should be in litellm_params"
                assert hasattr(
                    litellm_params, "aws_region_name"
                ), "aws_region_name should be in litellm_params"

                # Verify we got the expected response
                assert search_response["data"][0]["id"] == "test_result"

                print(
                    f"✅ Test passed: Credential accessor was called with vector store ID: T37J8R4WTM"
                )
                print(f"✅ Retrieved credentials: {retrieved_credentials}")
                print(f"✅ Credentials were injected into search call")
                print(f"✅ Search completed successfully using registry credentials")

    finally:
        # Restore original state
        litellm.vector_store_registry = original_registry
        litellm.credential_list = original_credential_list
