# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

# pylint: disable=protected-access

import time
import uuid
from typing import Dict, Iterable, Optional, cast

from marshmallow.exceptions import ValidationError as SchemaValidationError

from azure.ai.ml._exception_helper import log_and_raise_error
from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024Preview
from azure.ai.ml._restclient.v2024_01_01_preview.models import ComputeInstanceDataMount
from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview
from azure.ai.ml._restclient.v2024_07_01_preview.models import Datastore as DatastoreData
from azure.ai.ml._restclient.v2024_07_01_preview.models import DatastoreSecrets, NoneDatastoreCredentials, SecretExpiry
from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope, _ScopeDependentOperations
from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
from azure.ai.ml._utils._experimental import experimental
from azure.ai.ml._utils._logger_utils import OpsLogger
from azure.ai.ml.entities._datastore.datastore import Datastore
from azure.ai.ml.exceptions import MlException, ValidationException

ops_logger = OpsLogger(__name__)
module_logger = ops_logger.module_logger


class DatastoreOperations(_ScopeDependentOperations):
    """Represents a client for performing operations on Datastores.

    You should not instantiate this class directly. Instead, you should create MLClient and use this client via the
    property MLClient.datastores

    :param operation_scope: Scope variables for the operations classes of an MLClient object.
    :type operation_scope: ~azure.ai.ml._scope_dependent_operations.OperationScope
    :param operation_config: Common configuration for operations classes of an MLClient object.
    :type operation_config: ~azure.ai.ml._scope_dependent_operations.OperationConfig
    :param serviceclient_2024_01_01_preview: Service client to allow end users to operate on Azure Machine Learning
        Workspace resources.
    :type serviceclient_2024_01_01_preview: ~azure.ai.ml._restclient.v2023_01_01_preview.
        _azure_machine_learning_workspaces.AzureMachineLearningWorkspaces
    :param serviceclient_2024_07_01_preview: Service client to allow end users to operate on Azure Machine Learning
        Workspace resources.
    :type serviceclient_2024_07_01_preview: ~azure.ai.ml._restclient.v2024_07_01_preview.
        _azure_machine_learning_workspaces.AzureMachineLearningWorkspaces
    """

    def __init__(
        self,
        operation_scope: OperationScope,
        operation_config: OperationConfig,
        serviceclient_2024_01_01_preview: ServiceClient012024Preview,
        serviceclient_2024_07_01_preview: ServiceClient072024Preview,
        **kwargs: Dict,
    ):
        super(DatastoreOperations, self).__init__(operation_scope, operation_config)
        ops_logger.update_filter()
        self._operation = serviceclient_2024_07_01_preview.datastores
        self._compute_operation = serviceclient_2024_01_01_preview.compute
        self._credential = serviceclient_2024_07_01_preview._config.credential
        self._init_kwargs = kwargs

    @monitor_with_activity(ops_logger, "Datastore.List", ActivityType.PUBLICAPI)
    def list(self, *, include_secrets: bool = False) -> Iterable[Datastore]:
        """Lists all datastores and associated information within a workspace.

        :keyword include_secrets: Include datastore secrets in returned datastores, defaults to False
        :paramtype include_secrets: bool
        :return: An iterator like instance of Datastore objects
        :rtype: ~azure.core.paging.ItemPaged[Datastore]

        .. admonition:: Example:

            .. literalinclude:: ../samples/ml_samples_misc.py
                :start-after: [START datastore_operations_list]
                :end-before: [END datastore_operations_list]
                :language: python
                :dedent: 8
                :caption: List datastore example.
        """

        def _list_helper(datastore_resource: Datastore, include_secrets: bool) -> Datastore:
            if include_secrets:
                self._fetch_and_populate_secret(datastore_resource)
            return Datastore._from_rest_object(datastore_resource)

        return cast(
            Iterable[Datastore],
            self._operation.list(
                resource_group_name=self._operation_scope.resource_group_name,
                workspace_name=self._workspace_name,
                cls=lambda objs: [_list_helper(obj, include_secrets) for obj in objs],
                **self._init_kwargs,
            ),
        )

    @monitor_with_activity(ops_logger, "Datastore.ListSecrets", ActivityType.PUBLICAPI)
    def _list_secrets(self, name: str, expirable_secret: bool = False) -> DatastoreSecrets:
        return self._operation.list_secrets(
            name=name,
            body=SecretExpiry(expirable_secret=expirable_secret),
            resource_group_name=self._operation_scope.resource_group_name,
            workspace_name=self._workspace_name,
            **self._init_kwargs,
        )

    @monitor_with_activity(ops_logger, "Datastore.Delete", ActivityType.PUBLICAPI)
    def delete(self, name: str) -> None:
        """Deletes a datastore reference with the given name from the workspace. This method does not delete the actual
        datastore or underlying data in the datastore.

        :param name: Name of the datastore
        :type name: str

        .. admonition:: Example:

            .. literalinclude:: ../samples/ml_samples_misc.py
                :start-after: [START datastore_operations_delete]
                :end-before: [END datastore_operations_delete]
                :language: python
                :dedent: 8
                :caption: Delete datastore example.
        """

        self._operation.delete(
            name=name,
            resource_group_name=self._operation_scope.resource_group_name,
            workspace_name=self._workspace_name,
            **self._init_kwargs,
        )

    @monitor_with_activity(ops_logger, "Datastore.Get", ActivityType.PUBLICAPI)
    def get(self, name: str, *, include_secrets: bool = False) -> Datastore:  # type: ignore
        """Returns information about the datastore referenced by the given name.

        :param name: Datastore name
        :type name: str
        :keyword include_secrets: Include datastore secrets in the returned datastore, defaults to False
        :paramtype include_secrets: bool
        :return: Datastore with the specified name.
        :rtype: Datastore

        .. admonition:: Example:

            .. literalinclude:: ../samples/ml_samples_misc.py
                :start-after: [START datastore_operations_get]
                :end-before: [END datastore_operations_get]
                :language: python
                :dedent: 8
                :caption: Get datastore example.
        """
        try:
            datastore_resource = self._operation.get(
                name=name,
                resource_group_name=self._operation_scope.resource_group_name,
                workspace_name=self._workspace_name,
                **self._init_kwargs,
            )
            if include_secrets:
                self._fetch_and_populate_secret(datastore_resource)
            return Datastore._from_rest_object(datastore_resource)
        except (ValidationException, SchemaValidationError) as ex:
            log_and_raise_error(ex)

    def _fetch_and_populate_secret(self, datastore_resource: DatastoreData) -> None:
        if datastore_resource.name and not isinstance(
            datastore_resource.properties.credentials, NoneDatastoreCredentials
        ):
            secrets = self._list_secrets(name=datastore_resource.name, expirable_secret=True)
            datastore_resource.properties.credentials.secrets = secrets

    @monitor_with_activity(ops_logger, "Datastore.GetDefault", ActivityType.PUBLICAPI)
    def get_default(self, *, include_secrets: bool = False) -> Datastore:  # type: ignore
        """Returns the workspace's default datastore.

        :keyword include_secrets: Include datastore secrets in the returned datastore, defaults to False
        :paramtype include_secrets: bool
        :return: The default datastore.
        :rtype: Datastore

        .. admonition:: Example:

            .. literalinclude:: ../samples/ml_samples_misc.py
                :start-after: [START datastore_operations_get_default]
                :end-before: [END datastore_operations_get_default]
                :language: python
                :dedent: 8
                :caption: Get default datastore example.
        """
        try:
            datastore_resource = self._operation.list(
                resource_group_name=self._operation_scope.resource_group_name,
                workspace_name=self._workspace_name,
                is_default=True,
                **self._init_kwargs,
            ).next()
            if include_secrets:
                self._fetch_and_populate_secret(datastore_resource)
            return Datastore._from_rest_object(datastore_resource)
        except (ValidationException, SchemaValidationError) as ex:
            log_and_raise_error(ex)

    @monitor_with_activity(ops_logger, "Datastore.CreateOrUpdate", ActivityType.PUBLICAPI)
    def create_or_update(self, datastore: Datastore) -> Datastore:  # type: ignore
        """Attaches the passed in datastore to the workspace or updates the datastore if it already exists.

        :param datastore: The configuration of the datastore to attach.
        :type datastore: Datastore
        :return: The attached datastore.
        :rtype: Datastore

        .. admonition:: Example:

            .. literalinclude:: ../samples/ml_samples_misc.py
                :start-after: [START datastore_operations_create_or_update]
                :end-before: [END datastore_operations_create_or_update]
                :language: python
                :dedent: 8
                :caption: Create datastore example.
        """
        try:
            ds_request = datastore._to_rest_object()
            datastore_resource = self._operation.create_or_update(
                name=datastore.name,
                resource_group_name=self._operation_scope.resource_group_name,
                workspace_name=self._workspace_name,
                body=ds_request,
                skip_validation=True,
            )
            return Datastore._from_rest_object(datastore_resource)
        except Exception as ex:  # pylint: disable=W0718
            if isinstance(ex, (ValidationException, SchemaValidationError)):
                log_and_raise_error(ex)
            else:
                raise ex

    @monitor_with_activity(ops_logger, "Datastore.Mount", ActivityType.PUBLICAPI)
    @experimental
    def mount(
        self,
        path: str,
        *,
        mount_point: Optional[str] = None,
        mode: str = "ro_mount",
        debug: bool = False,
        persistent: bool = False,
        **kwargs,
    ) -> None:
        """Mount a datastore to a local path, so that you can access data inside it
        under a local path with any tools of your choice.

        :param path: The data store path to mount, in the form of `<name>` or `azureml://datastores/<name>`.
        :type path: str
        :keyword mount_point: A local path used as mount point.
        :type mount_point: str
        :keyword mode: Mount mode, either `ro_mount` (read-only) or `rw_mount` (read-write).
        :type mode: str
        :keyword debug: Whether to enable verbose logging.
        :type debug: bool
        :keyword persistent: Whether to persist the mount after reboot. Applies only when running on Compute Instance,
                where the 'CI_NAME' environment variable is set."
        :type persistent: bool
        :return: None
        """

        assert mode in ["ro_mount", "rw_mount"], "mode should be either `ro_mount` or `rw_mount`"
        read_only = mode == "ro_mount"

        import os

        ci_name = os.environ.get("CI_NAME")
        assert not persistent or (
            persistent and ci_name is not None
        ), "persistent mount is only supported on Compute Instance, where the 'CI_NAME' environment variable is set."

        try:
            from azureml.dataprep import rslex_fuse_subprocess_wrapper
        except ImportError as exc:
            msg = "Mount operations requires package azureml-dataprep-rslex installed. \
                You can install it with Azure ML SDK with `pip install azure-ai-ml[mount]`."
            raise MlException(message=msg, no_personal_data_message=msg) from exc

        uri = rslex_fuse_subprocess_wrapper.build_datastore_uri(
            self._operation_scope._subscription_id, self._resource_group_name, self._workspace_name, path
        )
        if persistent and ci_name is not None:
            mount_name = f"unified_mount_{str(uuid.uuid4()).replace('-', '')}"
            self._compute_operation.update_data_mounts(
                self._resource_group_name,
                self._workspace_name,
                ci_name,
                [
                    ComputeInstanceDataMount(
                        source=uri,
                        source_type="URI",
                        mount_name=mount_name,
                        mount_action="Mount",
                        mount_path=mount_point or "",
                    )
                ],
                api_version="2021-01-01",
                **kwargs,
            )
            print(f"Mount requested [name: {mount_name}]. Waiting for completion ...")
            while True:
                compute = self._compute_operation.get(self._resource_group_name, self._workspace_name, ci_name)
                mounts = compute.properties.properties.data_mounts
                try:
                    mount = [mount for mount in mounts if mount.mount_name == mount_name][0]
                    if mount.mount_state == "Mounted":
                        print(f"Mounted [name: {mount_name}].")
                        break
                    if mount.mount_state == "MountRequested":
                        pass
                    elif mount.mount_state == "MountFailed":
                        msg = f"Mount failed [name: {mount_name}]: {mount.error}"
                        raise MlException(message=msg, no_personal_data_message=msg)
                    else:
                        msg = f"Got unexpected mount state [name: {mount_name}]: {mount.mount_state}"
                        raise MlException(message=msg, no_personal_data_message=msg)
                except IndexError:
                    pass
                time.sleep(5)
        else:
            rslex_fuse_subprocess_wrapper.start_fuse_mount_subprocess(
                uri, mount_point, read_only, debug, credential=self._operation._config.credential
            )
