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

from msrest.serialization import Model


class DatasetRegistration(Model):
    """DatasetRegistration.

    :param name:
    :type name: str
    :param create_new_version:
    :type create_new_version: bool
    :param description:
    :type description: str
    :param tags:
    :type tags: dict[str, str]
    :param additional_transformations:
    :type additional_transformations: str
    """

    _attribute_map = {
        'name': {'key': 'name', 'type': 'str'},
        'create_new_version': {'key': 'createNewVersion', 'type': 'bool'},
        'description': {'key': 'description', 'type': 'str'},
        'tags': {'key': 'tags', 'type': '{str}'},
        'additional_transformations': {'key': 'additionalTransformations', 'type': 'str'},
    }

    def __init__(self, **kwargs):
        super(DatasetRegistration, self).__init__(**kwargs)
        self.name = kwargs.get('name', None)
        self.create_new_version = kwargs.get('create_new_version', None)
        self.description = kwargs.get('description', None)
        self.tags = kwargs.get('tags', None)
        self.additional_transformations = kwargs.get('additional_transformations', None)


class PipelineData(object):
    """
    Represents intermediate data in an Azure Machine Learning pipeline.

    Data used in pipeline can be produced by one step and consumed in another step by providing a PipelineData
    object as an output of one step and an input of one or more subsequent steps.

    Note that we do not expose this as part of our public API yet. This is only intended for internal usage.

    .. remarks::

        PipelineData represents data output a step will produce when it is run. Use PipelineData when creating steps
        to describe the files or directories which will be generated by the step. These data outputs will be added to
        the specified Datastore and can be retrieved and viewed later.

        For supported compute types, PipelineData can also be used to specify how the data will be produced and
        consumed by the run. There are two supported methods:

        * Mount (default): The input or output data is mounted to local storage on the compute
          node, and an environment variable is set which points to the path of this data ($AZUREML_DATAREFERENCE_name).
          For convenience, you can pass the PipelineData object in as one of the arguments to your script, for
          example using the ``arguments`` parameter of :class:`azureml.pipeline.steps.PythonScriptStep`, and
          the object will resolve to the path to the data. For outputs, your compute script should create a file
          or directory at this output path. To see the value of the environment variable used when you pass in the
          Pipeline object as an argument, use the :meth:`azureml.pipeline.core.PipelineData.get_env_variable_name`
          method.

        * Upload: Specify an ``output_path_on_compute`` corresponding to a file or directory
          name that your script will generate. (Environment variables are not used in this case.)


    :param name: The name of the PipelineData object, which can contain only letters, digits, and underscores.
    :type name: str
    :param datastore: The Datastore the PipelineData will reside on. If unspecified, the default datastore is used.
    :type datastore: azureml.data.azure_storage_datastore.AbstractAzureStorageDatastore or
        azureml.data.azure_data_lake_datastore.AzureDataLakeDatastore
    :param output_name: The name of the output, if None name is used. Can contain only letters, digits, and
        underscores.
    :type output_name: str
    :param output_mode: Specifies whether the producing step will use "upload" or "mount" method to access the data.
    :type output_mode: str
    :param output_path_on_compute: For ``output_mode`` = "upload", this parameter represents the path the
        module writes the output to.
    :type output_path_on_compute: str
    :param output_overwrite: For ``output_mode`` = "upload", this parameter specifies whether to overwrite
        existing data.
    :type output_overwrite: bool
    :param is_directory: Specifies whether the data is a directory or single file. This is only used to determine
        a data type used by Azure ML backend when the ``data_type`` parameter is not provided. The default is False.
    :type is_directory: bool
    :param dataset_registration: output dataset registration info.
    :type dataset_registration: DatasetRegistration
    """
    def __init__(self, name, datastore=None, output_name=None, output_mode="mount", output_path_on_compute=None,
                 output_overwrite=None, is_directory=None, dataset_registration: DatasetRegistration = None):
        import re

        if not name:
            raise ValueError("PipelineData name cannot be empty")

        invalid_name_exp = re.compile('\\W')
        if invalid_name_exp.search(name):
            raise ValueError("PipelineData name: [{name}] is not a valid, as it may contain only letters, "
                             "digits, and underscores.".format(name=name))

        self._name = name

        if output_name is None:
            output_name = name
        self._output_name = output_name

        import re
        invalid_name_exp = re.compile('\\W')
        if invalid_name_exp.search(output_name):
            raise ValueError("PipelineData output_name: [{name}] is not a valid, as it may contain only letters, "
                             "digits, and underscores.".format(name=output_name))

        self._datastore = datastore
        self._producer = None
        self._port_name = None
        self._output_mode = output_mode
        self._output_path_on_compute = output_path_on_compute
        self._output_overwrite = output_overwrite
        self._is_directory = is_directory
        self._dataset_registration = dataset_registration

        if self._output_mode not in ["mount", "upload"]:
            raise ValueError("Invalid output mode [%s]" % self._output_mode)

    def _set_producer(self, step):
        self._producer = step

    def _set_port_name(self, port_name):
        self._port_name = port_name

    @property
    def name(self):
        """
        Name of the PipelineData object.

        :return: Name.
        :rtype: str
        """
        return self._name

    @property
    def datastore(self):
        """
        Datastore the PipelineData will reside on.

        :return: The Datastore object.
        :rtype: azureml.data.azure_storage_datastore.AbstractAzureStorageDatastore or
            azureml.data.azure_data_lake_datastore.AzureDataLakeDatastore
        """
        return self._datastore

    def __str__(self):
        """
        __str__ override.

        :return: A string representation of the PipelineData.
        :rtype: str
        """
        return "$AZUREML_DATAREFERENCE_{0}".format(self.name)

    def __repr__(self):
        """
        Return __str__.

        :return: A string representation of the PipelineData.
        :rtype: str
        """
        return self.__str__()
