from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Final, Literal, Type, Union

from great_expectations._docs_decorators import public_api
from great_expectations.compatibility import azure, pydantic
from great_expectations.compatibility.typing_extensions import override
from great_expectations.core.util import AzureUrl
from great_expectations.datasource.fluent import _PandasFilePathDatasource
from great_expectations.datasource.fluent.config_str import (
    ConfigStr,
    _check_config_substitutions_needed,
)
from great_expectations.datasource.fluent.data_connector import (
    AzureBlobStorageDataConnector,
)
from great_expectations.datasource.fluent.interfaces import TestConnectionError
from great_expectations.datasource.fluent.pandas_datasource import PandasDatasourceError

if TYPE_CHECKING:
    from great_expectations.compatibility.azure import BlobServiceClient
    from great_expectations.datasource.fluent.data_asset.path.file_asset import FileDataAsset

logger = logging.getLogger(__name__)


_MISSING: Final = object()


class PandasAzureBlobStorageDatasourceError(PandasDatasourceError):
    pass


@public_api
class PandasAzureBlobStorageDatasource(_PandasFilePathDatasource):
    """
    PandasAzureBlobStorageDatasource is a PandasDatasource that uses Azure Blob Storage as a
    data store.
    """

    # class attributes
    data_connector_type: ClassVar[Type[AzureBlobStorageDataConnector]] = (
        AzureBlobStorageDataConnector
    )

    # instance attributes
    type: Literal["pandas_abs"] = "pandas_abs"

    # Azure Blob Storage specific attributes
    azure_options: Dict[str, Union[ConfigStr, Any]] = {}

    _account_name: str = pydantic.PrivateAttr(default="")
    # on 3.11 the annotation must be type-checking import otherwise it will fail at import time
    _azure_client: Union[BlobServiceClient, None] = pydantic.PrivateAttr(default=None)

    def _get_azure_client(self) -> azure.BlobServiceClient:
        azure_client: Union[azure.BlobServiceClient, None] = self._azure_client
        if not azure_client:
            _check_config_substitutions_needed(
                self, self.azure_options, raise_warning_if_provider_not_present=True
            )
            # pull in needed config substitutions using the `_config_provider`
            # The `FluentBaseModel.dict()` call will do the config substitution on the serialized dict if a `config_provider` is passed.  # noqa: E501 # FIXME CoP
            azure_options: dict = self.dict(config_provider=self._config_provider).get(
                "azure_options", {}
            )

            # Thanks to schema validation, we are guaranteed to have one of `conn_str` or `account_url` to  # noqa: E501 # FIXME CoP
            # use in authentication (but not both). If the format or content of the provided keys is invalid,  # noqa: E501 # FIXME CoP
            # the assignment of `self._account_name` and `self._azure_client` will fail and an error will be raised.  # noqa: E501 # FIXME CoP
            conn_str: str | None = azure_options.get("conn_str")
            account_url: str | None = azure_options.get("account_url")
            if not bool(conn_str) ^ bool(account_url):
                raise PandasAzureBlobStorageDatasourceError(  # noqa: TRY003 # FIXME CoP
                    "You must provide one of `conn_str` or `account_url` to the `azure_options` key in your config (but not both)"  # noqa: E501 # FIXME CoP
                )

            # Validate that "azure" libararies were successfully imported and attempt to create "azure_client" handle.  # noqa: E501 # FIXME CoP
            if azure.BlobServiceClient:  # type: ignore[truthy-function] # False if NotImported
                try:
                    if conn_str is not None:
                        self._account_name = re.search(  # type: ignore[union-attr] # FIXME CoP
                            r".*?AccountName=(.+?);.*?", conn_str
                        ).group(1)
                        azure_client = azure.BlobServiceClient.from_connection_string(
                            **azure_options
                        )
                    elif account_url is not None:
                        self._account_name = re.search(  # type: ignore[union-attr] # FIXME CoP
                            r"(?:https?://)?(.+?).blob.core.windows.net", account_url
                        ).group(1)
                        azure_client = azure.BlobServiceClient(**azure_options)
                except Exception as e:
                    # Failure to create "azure_client" is most likely due invalid "azure_options" dictionary.  # noqa: E501 # FIXME CoP
                    raise PandasAzureBlobStorageDatasourceError(  # noqa: TRY003 # FIXME CoP
                        f'Due to exception: "{e!s}", "azure_client" could not be created.'
                    ) from e
            else:
                raise PandasAzureBlobStorageDatasourceError(  # noqa: TRY003 # FIXME CoP
                    'Unable to create "PandasAzureBlobStorageDatasource" due to missing azure.storage.blob dependency.'  # noqa: E501 # FIXME CoP
                )

            self._azure_client = azure_client

        if not azure_client:
            raise PandasAzureBlobStorageDatasourceError("Failed to return `azure_client`")  # noqa: TRY003 # FIXME CoP

        return azure_client

    @override
    def test_connection(self, test_assets: bool = True) -> None:
        """Test the connection for the PandasAzureBlobStorageDatasource.

        Args:
            test_assets: If assets have been passed to the PandasAzureBlobStorageDatasource, whether to test them as well.

        Raises:
            TestConnectionError: If the connection test fails.
        """  # noqa: E501 # FIXME CoP
        try:
            _ = self._get_azure_client()
        except Exception as e:
            raise TestConnectionError(  # noqa: TRY003 # FIXME CoP
                f"Attempt to connect to datasource failed with the following error message: {e!s}"
            ) from e

        if self.assets and test_assets:
            for asset in self.assets:
                asset.test_connection()

    @override
    def _build_data_connector(
        self,
        data_asset: FileDataAsset,
        abs_container: str = _MISSING,  # type: ignore[assignment] # _MISSING is used as sentinel value
        abs_name_starts_with: str = "",
        abs_delimiter: str = "/",
        abs_recursive_file_discovery: bool = False,
        **kwargs,
    ) -> None:
        """Builds and attaches the `AzureBlobStorageDataConnector` to the asset."""
        if kwargs:
            raise TypeError(  # noqa: TRY003 # FIXME CoP
                f"_build_data_connector() got unexpected keyword arguments {list(kwargs.keys())}"
            )
        if abs_container is _MISSING:
            raise TypeError(f"'{data_asset.name}' is missing required argument 'abs_container'")  # noqa: TRY003 # FIXME CoP

        data_asset._data_connector = self.data_connector_type.build_data_connector(
            datasource_name=self.name,
            data_asset_name=data_asset.name,
            azure_client=self._get_azure_client(),
            account_name=self._account_name,
            container=abs_container,
            name_starts_with=abs_name_starts_with,
            delimiter=abs_delimiter,
            recursive_file_discovery=abs_recursive_file_discovery,
            file_path_template_map_fn=AzureUrl.AZURE_BLOB_STORAGE_HTTPS_URL_TEMPLATE.format,
        )

        # build a more specific `_test_connection_error_message`
        data_asset._test_connection_error_message = (
            self.data_connector_type.build_test_connection_error_message(
                data_asset_name=data_asset.name,
                account_name=self._account_name,
                container=abs_container,
                name_starts_with=abs_name_starts_with,
                delimiter=abs_delimiter,
                recursive_file_discovery=abs_recursive_file_discovery,
            )
        )
