from collections import defaultdict
from unittest.mock import AsyncMock, Mock, patch

import pytest

# Assuming the component is imported from the module
# Adjust the import path as needed
from lfx.base.datastax import AstraDBBaseComponent


@pytest.fixture
def mock_component():
    """Create a mock AstraDBBaseComponent instance."""
    component = AstraDBBaseComponent()
    component.token = "test_token"  # noqa: S105
    component.environment = "prod"
    component.database_name = "test_db"
    component.api_endpoint = None
    component.keyspace = None
    component.collection_name = "test_collection"
    component.log = Mock()
    return component


@pytest.fixture
def mock_database_info():
    """Mock database information structure."""
    return {
        "test_db": {
            "api_endpoints": ["https://test-db-id.apps.astra.datastax.com"],
            "keyspaces": ["default_keyspace", "custom_keyspace"],
            "collections": 5,
            "status": "ACTIVE",
            "org_id": "test-org-id",
        },
        "pending_db": {
            "api_endpoints": ["https://pending-db-id.apps.astra.datastax.com"],
            "keyspaces": ["default_keyspace"],
            "collections": 0,
            "status": "PENDING",
            "org_id": "test-org-id",
        },
    }


@pytest.fixture
def mock_build_config():
    """Mock build configuration structure."""
    return {
        "token": {"value": "test_token"},
        "environment": {"value": "prod"},
        "database_name": {
            "value": "test_db",
            "options": [],
            "options_metadata": [],
            "show": False,
            "dialog_inputs": {
                "fields": {
                    "data": {
                        "node": {
                            "template": {
                                "02_cloud_provider": {"options": []},
                                "03_region": {"options": [], "value": None},
                            }
                        }
                    }
                }
            },
        },
        "api_endpoint": {"value": "", "options": []},
        "keyspace": {"value": "", "options": []},
        "collection_name": {
            "value": "",
            "options": [],
            "options_metadata": [],
            "show": False,
            "dialog_inputs": {
                "fields": {
                    "data": {
                        "node": {
                            "template": {
                                "01_new_collection_name": {},
                                "02_embedding_generation_provider": {
                                    "value": None,
                                    "options": [],
                                    "options_metadata": [],
                                    "helper_text": "",
                                },
                                "03_embedding_generation_model": {
                                    "options": [],
                                    "placeholder": None,
                                    "readonly": False,
                                    "required": True,
                                    "value": None,
                                },
                                "04_dimension": {
                                    "placeholder": None,
                                    "value": None,
                                    "readonly": False,
                                    "required": False,
                                },
                            }
                        }
                    }
                }
            },
        },
        "embedding_model": {"show": True, "required": False},
        "autodetect_collection": {"value": True},
        "search_method": {"value": "Vector Search", "show": False, "options": []},
        "search_type": {"value": "Similarity", "show": True},
        "search_score_threshold": {"show": True},
        "reranker": {
            "value": "",
            "options": [],
            "options_metadata": [],
            "show": False,
            "toggle_disable": False,
            "toggle_value": True,
        },
        "lexical_terms": {"value": "", "show": False},
    }


class TestCloudProviderMapping:
    """Tests for cloud provider mapping."""

    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    def test_map_cloud_providers_structure(self, mock_client_class):
        """Test that map_cloud_providers returns correct structure."""
        # Mock the admin client and its methods
        mock_admin = Mock()
        # Mock find_available_regions to return a list of region objects
        mock_region1 = Mock()
        mock_region1.cloud_provider = "AWS"
        mock_region1.name = "us-east-2"
        mock_region2 = Mock()
        mock_region2.cloud_provider = "GCP"
        mock_region2.name = "us-central1"
        mock_region3 = Mock()
        mock_region3.cloud_provider = "Azure"
        mock_region3.name = "eastus"

        mock_admin.find_available_regions.return_value = [mock_region1, mock_region2, mock_region3]

        mock_client = mock_client_class.return_value
        mock_client.get_admin.return_value = mock_admin

        providers = AstraDBBaseComponent.map_cloud_providers(token="test_token")  # noqa: S106

        assert "Amazon Web Services" in providers
        assert "Google Cloud Platform" in providers
        assert "Microsoft Azure" in providers

    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    def test_map_cloud_providers_prod_content(self, mock_client_class):
        """Test production environment cloud providers."""
        # Mock the admin client and its methods
        mock_admin = Mock()
        # Mock find_available_regions to return a list of region objects
        mock_region1 = Mock()
        mock_region1.cloud_provider = "AWS"
        mock_region1.name = "us-east-2"
        mock_region2 = Mock()
        mock_region2.cloud_provider = "AWS"
        mock_region2.name = "us-west-2"
        mock_region3 = Mock()
        mock_region3.cloud_provider = "GCP"
        mock_region3.name = "us-central1"
        mock_region4 = Mock()
        mock_region4.cloud_provider = "Azure"
        mock_region4.name = "eastus"

        mock_admin.find_available_regions.return_value = [mock_region1, mock_region2, mock_region3, mock_region4]

        mock_client = mock_client_class.return_value
        mock_client.get_admin.return_value = mock_admin

        providers = AstraDBBaseComponent.map_cloud_providers(token="test_token")  # noqa: S106

        assert "Amazon Web Services" in providers
        assert "Google Cloud Platform" in providers
        assert "Microsoft Azure" in providers

        assert providers["Amazon Web Services"]["id"] == "aws"
        assert "us-east-2" in providers["Amazon Web Services"]["regions"]
        assert "us-west-2" in providers["Amazon Web Services"]["regions"]

    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    def test_map_cloud_providers_dev_content(self, mock_client_class):
        """Test development environment cloud providers."""
        # Mock the admin client and its methods
        mock_admin = Mock()
        # Mock find_available_regions to return a list of region objects
        mock_region1 = Mock()
        mock_region1.cloud_provider = "AWS"
        mock_region1.name = "us-east-2"
        mock_region2 = Mock()
        mock_region2.cloud_provider = "GCP"
        mock_region2.name = "us-central1"

        mock_admin.find_available_regions.return_value = [mock_region1, mock_region2]

        mock_client = mock_client_class.return_value
        mock_client.get_admin.return_value = mock_admin

        providers = AstraDBBaseComponent.map_cloud_providers(token="test_token")  # noqa: S106

        assert "Amazon Web Services" in providers
        assert "Google Cloud Platform" in providers
        assert providers["Amazon Web Services"]["id"] == "aws"


class TestDatabaseIdExtraction:
    """Tests for database ID extraction from API endpoints."""

    def test_get_database_id_static_valid_uuid(self):
        """Test extracting valid UUID from API endpoint."""
        api_endpoint = "https://12345678-1234-1234-1234-123456789abc.apps.astra.datastax.com"
        db_id = AstraDBBaseComponent.get_database_id_static(api_endpoint)

        assert db_id == "12345678-1234-1234-1234-123456789abc"

    def test_get_database_id_static_no_uuid(self):
        """Test extraction returns None when no UUID present."""
        api_endpoint = "https://invalid.endpoint.com"
        db_id = AstraDBBaseComponent.get_database_id_static(api_endpoint)

        assert db_id is None

    def test_get_database_id_static_case_insensitive(self):
        """Test UUID extraction is case insensitive."""
        api_endpoint = "https://ABCDEF12-ABCD-ABCD-ABCD-ABCDEFABCDEF.apps.astra.datastax.com"
        db_id = AstraDBBaseComponent.get_database_id_static(api_endpoint)

        assert db_id == "ABCDEF12-ABCD-ABCD-ABCD-ABCDEFABCDEF"

    def test_get_database_id_instance_method(self, mock_component):
        """Test instance method for getting database ID."""
        with patch.object(
            mock_component,
            "get_api_endpoint",
            return_value="https://12345678-1234-1234-1234-123456789abc.apps.astra.datastax.com",
        ):
            db_id = mock_component.get_database_id()
            assert db_id == "12345678-1234-1234-1234-123456789abc"


class TestKeyspaceHandling:
    """Tests for keyspace handling."""

    def test_get_keyspace_default(self, mock_component):
        """Test getting default keyspace when none set."""
        mock_component.keyspace = None
        assert mock_component.get_keyspace() == "default_keyspace"

    def test_get_keyspace_with_value(self, mock_component):
        """Test getting keyspace when value is set."""
        mock_component.keyspace = "custom_keyspace"
        assert mock_component.get_keyspace() == "custom_keyspace"

    def test_get_keyspace_with_empty_string(self, mock_component):
        """Test getting keyspace with empty string returns default."""
        mock_component.keyspace = ""
        assert mock_component.get_keyspace() == "default_keyspace"


class TestCollectionDataRetrieval:
    """Tests for collection data retrieval."""

    def test_collection_data_success(self, mock_component):
        """Test successful collection data retrieval."""
        mock_database = Mock()
        mock_collection = Mock()
        mock_collection.estimated_document_count.return_value = 100
        mock_database.get_collection.return_value = mock_collection

        count = mock_component.collection_data("test_collection", database=mock_database)
        assert count == 100
        mock_database.get_collection.assert_called_once_with("test_collection")
        mock_collection.estimated_document_count.assert_called_once()

    def test_collection_data_with_provided_database(self, mock_component):
        """Test collection data retrieval with provided database object."""
        mock_database = Mock()
        mock_collection = Mock()
        mock_collection.estimated_document_count.return_value = 50

        mock_database.get_collection.return_value = mock_collection

        count = mock_component.collection_data("test_collection", database=mock_database)
        assert count == 50

    def test_collection_data_error_handling(self, mock_component):
        """Test collection data error handling."""
        mock_database = Mock()
        mock_database.get_collection.side_effect = Exception("Connection error")

        result = mock_component.collection_data("test_collection", database=mock_database)
        assert result is None
        mock_component.log.assert_called_once()


class TestDatabaseCreation:
    """Tests for database creation."""

    @pytest.mark.asyncio
    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    async def test_create_database_api_success(self, mock_client_class):
        """Test successful database creation."""
        mock_admin = Mock()
        # Fix: Make async_create_database return a proper awaitable that yields a dict
        mock_admin.async_create_database = AsyncMock(return_value={"id": "new-db-id"})

        # Mock find_available_regions to return a list of region objects
        mock_region = Mock()
        mock_region.cloud_provider = "AWS"
        mock_region.name = "us-east-2"
        mock_admin.find_available_regions.return_value = [mock_region]

        mock_client = mock_client_class.return_value
        mock_client.get_admin.return_value = mock_admin

        result = await AstraDBBaseComponent.create_database_api(
            new_database_name="new_db",
            cloud_provider="Amazon Web Services",
            region="us-east-2",
            token="test_token",  # noqa: S106
            environment="prod",
        )

        assert result == {"id": "new-db-id"}
        mock_admin.async_create_database.assert_called_once()


class TestCollectionCreation:
    """Tests for collection creation."""

    @pytest.mark.asyncio
    @patch("lfx.base.datastax.astradb_base._AstraDBCollectionEnvironment")
    @patch.object(AstraDBBaseComponent, "get_vectorize_providers")
    async def test_create_collection_api_with_vectorize(self, mock_get_providers, mock_env_class):
        """Test collection creation with vectorize options."""
        mock_get_providers.return_value = defaultdict(
            list,
            {
                "NVIDIA": ["nvidia", ["model1", "model2"]],
            },
        )

        await AstraDBBaseComponent.create_collection_api(
            new_collection_name="new_collection",
            token="test_token",  # noqa: S106
            api_endpoint="https://test.endpoint.com",
            keyspace="default_keyspace",
            embedding_generation_provider="NVIDIA",
            embedding_generation_model="model1",
        )

        mock_env_class.assert_called_once()
        call_kwargs = mock_env_class.call_args[1]
        assert call_kwargs["collection_name"] == "new_collection"
        assert call_kwargs["collection_vector_service_options"] is not None

    @pytest.mark.asyncio
    @patch("lfx.base.datastax.astradb_base._AstraDBCollectionEnvironment")
    async def test_create_collection_api_with_dimension(self, mock_env_class):
        """Test collection creation with explicit dimension."""
        await AstraDBBaseComponent.create_collection_api(
            new_collection_name="new_collection",
            token="test_token",  # noqa: S106
            api_endpoint="https://test.endpoint.com",
            dimension=1536,
            keyspace="default_keyspace",
        )

        mock_env_class.assert_called_once()
        call_kwargs = mock_env_class.call_args[1]
        assert call_kwargs["embedding_dimension"] == 1536
        assert call_kwargs["collection_vector_service_options"] is None

    @pytest.mark.asyncio
    async def test_create_collection_api_no_name(self):
        """Test collection creation fails without name."""
        with pytest.raises(ValueError, match="Collection name is required"):
            await AstraDBBaseComponent.create_collection_api(
                new_collection_name="",
                token="test_token",  # noqa: S106
                api_endpoint="https://test.endpoint.com",
                keyspace="default_keyspace",
            )


class TestUpdateBuildConfig:
    """Tests for update_build_config method."""

    @pytest.mark.asyncio
    async def test_update_build_config_no_token(self, mock_component, mock_build_config):
        """Test update_build_config with no token resets config."""
        mock_component.token = None
        result = await mock_component.update_build_config(mock_build_config, "", "database_name")

        assert result["database_name"]["options"] == []
        assert result["database_name"]["show"] is False

    @pytest.mark.asyncio
    @patch.object(AstraDBBaseComponent, "_initialize_database_options")
    @patch.object(AstraDBBaseComponent, "map_cloud_providers")
    async def test_update_build_config_first_run(
        self, mock_map_providers, mock_init_db, mock_component, mock_build_config
    ):
        """Test update_build_config on first run."""
        # Mock the cloud providers mapping to avoid API calls
        mock_map_providers.return_value = {
            "prod": {
                "Amazon Web Services": {
                    "id": "aws",
                    "regions": ["us-east-2", "us-west-2"],
                }
            },
            "dev": {},
            "test": {},
        }

        mock_init_db.return_value = [
            {
                "name": "db1",
                "status": None,
                "collections": 5,
                "api_endpoints": ["https://db1.endpoint.com"],
                "keyspaces": ["default_keyspace"],
                "org_id": "org-id",
            }
        ]

        result = await mock_component.update_build_config(mock_build_config, "", "embedding_model")

        assert "db1" in result["database_name"]["options"]
        mock_init_db.assert_called_once()

    @pytest.mark.asyncio
    async def test_update_build_config_search_method_change_to_hybrid(self, mock_component, mock_build_config):
        """Test base update_build_config doesn't handle search_method."""
        # The base AstraDBBaseComponent doesn't handle search_method changes
        # This functionality is in the AstraDBVectorStoreComponent subclass
        result = await mock_component.update_build_config(mock_build_config, "Hybrid Search", "search_method")

        # Base component should return config unchanged for search_method
        assert result["lexical_terms"]["show"] is False  # Default value from fixture
        assert result["reranker"]["show"] is False  # Default value from fixture


class TestGetDatabaseObject:
    """Tests for getting database object."""

    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    def test_get_database_object_success(self, mock_client_class, mock_component):
        """Test successful database object retrieval."""
        mock_database = Mock()
        mock_client = mock_client_class.return_value
        mock_client.get_database.return_value = mock_database

        with (
            patch.object(mock_component, "get_api_endpoint", return_value="https://test.endpoint.com"),
            patch.object(mock_component, "get_keyspace", return_value="default_keyspace"),
        ):
            db = mock_component.get_database_object()
            assert db == mock_database

    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    def test_get_database_object_with_custom_endpoint(self, mock_client_class, mock_component):
        """Test database object retrieval with custom endpoint."""
        mock_database = Mock()
        mock_client = mock_client_class.return_value
        mock_client.get_database.return_value = mock_database

        with patch.object(mock_component, "get_keyspace", return_value="default_keyspace"):
            db = mock_component.get_database_object(api_endpoint="https://custom.endpoint.com")
            assert db == mock_database

    @patch("lfx.base.datastax.astradb_base.DataAPIClient")
    def test_get_database_object_error(self, mock_client_class, mock_component):
        """Test database object retrieval error handling."""
        mock_client_class.side_effect = Exception("Connection error")

        with (
            patch.object(mock_component, "get_api_endpoint", return_value="https://test.endpoint.com"),
            pytest.raises(ValueError, match="Error fetching database object"),
        ):
            mock_component.get_database_object()


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
