# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------
from inspect import Parameter
from typing import Dict, List, Optional, Union

from azure.ai.ml.constants._component import ExternalDataType
from azure.ai.ml.entities._inputs_outputs.utils import _remove_empty_values
from azure.ai.ml.entities._mixins import DictMixin, RestTranslatableMixin


class StoredProcedureParameter(DictMixin, RestTranslatableMixin):
    """Define a stored procedure parameter class for DataTransfer import database task.

    :keyword name: The name of the database stored procedure.
    :paramtype name: str
    :keyword value: The value of the database stored procedure.
    :paramtype value: str
    :keyword type: The type of the database stored procedure.
    :paramtype type: str
    """

    def __init__(
        self,
        *,
        name: Optional[str] = None,
        value: Optional[str] = None,
        type: Optional[str] = None,  # pylint: disable=redefined-builtin
    ) -> None:
        self.type = type
        self.name = name
        self.value = value


class Database(DictMixin, RestTranslatableMixin):  # pylint: disable=too-many-instance-attributes
    """Define a database class for a DataTransfer Component or Job.

    :keyword query: The SQL query to retrieve data from the database.
    :paramtype query: str
    :keyword table_name: The name of the database table.
    :paramtype table_name: str
    :keyword stored_procedure: The name of the stored procedure.
    :paramtype stored_procedure: str
    :keyword stored_procedure_params: The parameters for the stored procedure.
    :paramtype stored_procedure_params: List[dict, StoredProcedureParameter]
    :keyword connection: The connection string for the database.
        The credential information should be stored in the workspace connection.
    :paramtype connection: str
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if the Database object cannot be successfully validated.
        Details will be provided in the error message.

    .. admonition:: Example:

        .. literalinclude:: ../../../../../samples/ml_samples_input_output_configurations.py
            :start-after: [START configure_database]
            :end-before: [END configure_database]
            :language: python
            :dedent: 8
            :caption: Create a database and querying a database table.
    """

    _EMPTY = Parameter.empty

    def __init__(
        self,
        *,
        query: Optional[str] = None,
        table_name: Optional[str] = None,
        stored_procedure: Optional[str] = None,
        stored_procedure_params: Optional[List[dict]] = None,
        connection: Optional[str] = None,
    ) -> None:
        # As an annotation, it is not allowed to initialize the name.
        # The name will be updated by the annotated variable name.
        self.name = None
        self.type = ExternalDataType.DATABASE
        self.connection = connection
        self.query = query
        self.table_name = table_name
        self.stored_procedure = stored_procedure
        self.stored_procedure_params = stored_procedure_params

    def _to_dict(self, remove_name: bool = True) -> Dict:
        """Convert the Source object to a dict.

        :param remove_name: Whether to remove the `name` key from the  dict representation. Defaults to True.
        :type remove_name: bool
        :return: The dictionary representation of the class
        :rtype: Dict
        """
        keys = [
            "name",
            "type",
            "query",
            "stored_procedure",
            "stored_procedure_params",
            "connection",
            "table_name",
        ]
        if remove_name:
            keys.remove("name")
        result = {key: getattr(self, key) for key in keys}
        return _remove_empty_values(result)

    def _to_rest_object(self) -> Dict:
        # this is for component rest object when using Source as component inputs, as for job input usage,
        # rest object is generated by extracting Source's properties, see details in to_rest_dataset_literal_inputs()
        result = self._to_dict()
        return result

    def _update_name(self, name):
        self.name = name

    @classmethod
    def _from_rest_object(cls, obj: Dict) -> "Database":
        return Database(**obj)

    @property
    def stored_procedure_params(self) -> Optional[StoredProcedureParameter]:
        """Get or set the parameters for the stored procedure.

        :return: The parameters for the stored procedure.
        :rtype: List[StoredProcedureParameter]
        """

        return self._stored_procedure_params

    @stored_procedure_params.setter
    def stored_procedure_params(self, value: Union[Dict[str, str], StoredProcedureParameter, None]):
        """Set the parameters for the stored procedure.

        :param value: The parameters for the stored procedure.
        :type value: Union[Dict[str, str], StoredProcedureParameter, None]
        """
        if value is None:
            self._stored_procedure_params = value
        else:
            if not isinstance(value, list):
                value = [value]
            for index, item in enumerate(value):
                if isinstance(item, dict):
                    value[index] = StoredProcedureParameter(**item)
            self._stored_procedure_params = value


class FileSystem(DictMixin, RestTranslatableMixin):  # pylint: disable=too-many-instance-attributes
    """Define a file system class of a DataTransfer Component or Job.

    e.g. source_s3 = FileSystem(path='s3://my_bucket/my_folder', connection='azureml:my_s3_connection')

    :param path: The path to which the input is pointing. Could be pointing to the path of file system. Default is None.
    :type path: str
    :param connection: Connection is workspace, we didn't support storage connection here, need leverage workspace
        connection to store these credential info. Default is None.
    :type connection: str
    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if Source cannot be successfully validated.
        Details will be provided in the error message.
    """

    _EMPTY = Parameter.empty

    def __init__(
        self,
        *,
        path: Optional[str] = None,
        connection: Optional[str] = None,
    ) -> None:
        self.type = ExternalDataType.FILE_SYSTEM
        self.name = None
        self.connection = connection

        if path is not None and not isinstance(path, str):
            # this logic will make dsl data binding expression working in the same way as yaml
            # it's written to handle InputOutputBase, but there will be loop import if we import InputOutputBase here
            self.path = str(path)
        else:
            self.path = path

    def _to_dict(self, remove_name: bool = True):
        """Convert the Source object to a dict.

        :param remove_name: Whether to remove the `name` key from the  dict representation. Defaults to True.
        :type remove_name: bool
        :return: The dictionary representation of the object
        :rtype: Dict
        """
        keys = ["name", "path", "type", "connection"]
        if remove_name:
            keys.remove("name")
        result = {key: getattr(self, key) for key in keys}
        return _remove_empty_values(result)

    def _to_rest_object(self) -> Dict:
        # this is for component rest object when using Source as component inputs, as for job input usage,
        # rest object is generated by extracting Source's properties, see details in to_rest_dataset_literal_inputs()
        result = self._to_dict()
        return result

    def _update_name(self, name):
        self.name = name

    @classmethod
    def _from_rest_object(cls, obj: Dict) -> "FileSystem":
        return FileSystem(**obj)
