# -*- coding: utf-8 -*-

import copy
import logging

from ruamel.yaml import YAML

import warnings

from great_expectations.data_context.util import verify_dynamic_loading_support
from great_expectations.data_context.util import (
    load_class,
    instantiate_class_from_config
)
from great_expectations.types import ClassConfig
from great_expectations.exceptions import ClassInstantiationError

logger = logging.getLogger(__name__)
yaml = YAML()
yaml.default_flow_style = False


class Datasource(object):
    """A Datasource connects to a compute environment and one or more storage environments and produces batches of data
    that Great Expectations can validate in that compute environment.

    Each Datasource provides Batches connected to a specific compute environment, such as a
    SQL database, a Spark cluster, or a local in-memory Pandas DataFrame.

    Datasources use Batch Kwargs to specify instructions for how to access data from
    relevant sources such as an existing object from a DAG runner, a SQL database, S3 bucket, or local filesystem.

    To bridge the gap between those worlds, Datasources interact closely with *generators* which
    are aware of a source of data and can produce produce identifying information, called
    "batch_kwargs" that datasources can use to get individual batches of data. They add flexibility
    in how to obtain data such as with time-based partitioning, downsampling, or other techniques
    appropriate for the datasource.

    For example, a batch kwargs generator could produce a SQL query that logically represents "rows in the Events
    table with a timestamp on February 7, 2012," which a SqlAlchemyDatasource could use to materialize
    a SqlAlchemyDataset corresponding to that batch of data and ready for validation.

    Since opinionated DAG managers such as airflow, dbt, prefect.io, dagster can also act as datasources
    and/or batch kwargs generators for a more generic datasource.

    When adding custom expectations by subclassing an existing DataAsset type, use the data_asset_type parameter
    to configure the datasource to load and return DataAssets of the custom type.
    """
    recognized_batch_parameters = {'limit'}

    @classmethod
    def from_configuration(cls, **kwargs):
        """
        Build a new datasource from a configuration dictionary.

        Args:
            **kwargs: configuration key-value pairs

        Returns:
            datasource (Datasource): the newly-created datasource

        """
        return cls(**kwargs)

    @classmethod
    def build_configuration(
            cls,
            class_name,
            module_name="great_expectations.datasource",
            data_asset_type=None,
            batch_kwargs_generators=None,
            **kwargs
    ):
        """
        Build a full configuration object for a datasource, potentially including batch kwargs generators with defaults.

        Args:
            class_name: The name of the class for which to build the config
            module_name: The name of the module in which the datasource class is located
            data_asset_type: A ClassConfig dictionary
            batch_kwargs_generators: BatchKwargGenerators configuration dictionary
            **kwargs: Additional kwargs to be part of the datasource constructor's initialization

        Returns:
            A complete datasource configuration.

        """
        verify_dynamic_loading_support(module_name=module_name)
        class_ = load_class(class_name=class_name, module_name=module_name)
        configuration = class_.build_configuration(data_asset_type=data_asset_type, batch_kwargs_generators=batch_kwargs_generators, **kwargs)
        return configuration

    def __init__(self, name, data_context=None, data_asset_type=None, batch_kwargs_generators=None, **kwargs):
        """
        Build a new datasource.

        Args:
            name: the name for the datasource
            data_context: data context to which to connect
            data_asset_type (ClassConfig): the type of DataAsset to produce
            batch_kwargs_generators: BatchKwargGenerators to add to the datasource
        """
        self._data_context = data_context
        self._name = name
        if isinstance(data_asset_type, str):
            warnings.warn(
                "String-only configuration for data_asset_type is deprecated. Use module_name and class_name instead.",
                DeprecationWarning)
        self._data_asset_type = data_asset_type
        self._datasource_config = kwargs
        self._batch_kwargs_generators = {}

        self._datasource_config["data_asset_type"] = data_asset_type
        if batch_kwargs_generators is not None:
            self._datasource_config["batch_kwargs_generators"] = batch_kwargs_generators

    @property
    def name(self):
        """
        Property for datasource name
        """
        return self._name

    @property
    def config(self):
        return copy.deepcopy(self._datasource_config)

    @property
    def data_context(self):
        """
        Property for attached DataContext
        """
        return self._data_context

    def _build_generators(self):
        """
        Build batch kwargs generator objects from the datasource configuration.

        Returns:
            None
        """
        try:
            for generator in self._datasource_config["batch_kwargs_generators"].keys():
                self.get_batch_kwargs_generator(generator)
        except KeyError:
            pass

    def add_batch_kwargs_generator(self, name, class_name, **kwargs):
        """Add a BatchKwargGenerator to the datasource.

        Args:
            name (str): the name of the new BatchKwargGenerator to add
            class_name: class of the BatchKwargGenerator to add
            kwargs: additional keyword arguments will be passed directly to the new BatchKwargGenerator's constructor

        Returns:
             BatchKwargGenerator (BatchKwargGenerator)
        """
        kwargs["class_name"] = class_name
        generator = self._build_batch_kwargs_generator(**kwargs)
        if "batch_kwargs_generators" not in self._datasource_config:
            self._datasource_config["batch_kwargs_generators"] = dict()
        self._datasource_config["batch_kwargs_generators"][name] = kwargs

        return generator

    def _build_batch_kwargs_generator(self, **kwargs):
        """Build a BatchKwargGenerator using the provided configuration and return the newly-built generator."""
        generator = instantiate_class_from_config(
            config=kwargs,
            runtime_environment={
                "datasource": self
            },
            config_defaults={
                "module_name": "great_expectations.datasource.batch_kwargs_generator"
            }
        )
        if not generator:
            raise ClassInstantiationError(
                module_name="great_expectations.datasource.batch_kwargs_generator",
                package_name=None,
                class_name=kwargs['class_name']
            )

        return generator

    def get_batch_kwargs_generator(self, name):
        """Get the (named) BatchKwargGenerator from a datasource)

        Args:
            name (str): name of BatchKwargGenerator (default value is 'default')

        Returns:
            BatchKwargGenerator (BatchKwargGenerator)
        """
        if name in self._batch_kwargs_generators:
            return self._batch_kwargs_generators[name]
        elif "batch_kwargs_generators" in self._datasource_config and name in self._datasource_config["batch_kwargs_generators"]:
            generator_config = copy.deepcopy(self._datasource_config["batch_kwargs_generators"][name])
        else:
            raise ValueError(
                "Unable to load batch kwargs generator %s -- no configuration found or invalid configuration." % name
            )
        generator = self._build_batch_kwargs_generator(**generator_config)
        self._batch_kwargs_generators[name] = generator
        return generator

    def list_batch_kwargs_generators(self):
        """List currently-configured BatchKwargGenerator for this datasource.

        Returns:
            List(dict): each dictionary includes "name" and "type" keys
        """
        generators = []

        if "batch_kwargs_generators" in self._datasource_config:
            for key, value in self._datasource_config["batch_kwargs_generators"].items():
                generators.append({
                    "name": key,
                    "class_name": value["class_name"]
                })

        return generators

    def process_batch_parameters(self, limit=None, dataset_options=None):
        """Use datasource-specific configuration to translate any batch parameters into batch kwargs at the datasource
        level.

        Args:
            limit (int): a parameter all datasources must accept to allow limiting a batch to a smaller number of rows.
            dataset_options (dict): a set of kwargs that will be passed to the constructor of a dataset built using
                these batch_kwargs

        Returns:
            batch_kwargs: Result will include both parameters passed via argument and configured parameters.
        """
        batch_kwargs = self._datasource_config.get("batch_kwargs", {})

        if limit is not None:
            batch_kwargs["limit"] = limit

        if dataset_options is not None:
            # Then update with any locally-specified reader options
            if not batch_kwargs.get("dataset_options"):
                batch_kwargs["dataset_options"] = dict()
            batch_kwargs["dataset_options"].update(dataset_options)

        return batch_kwargs

    def get_batch(self, batch_kwargs, batch_parameters=None):
        """Get a batch of data from the datasource.

        Args:
            batch_kwargs: the BatchKwargs to use to construct the batch
            batch_parameters: optional parameters to store as the reference description of the batch. They should
                reflect parameters that would provide the passed BatchKwargs.


        Returns:
            Batch

        """
        raise NotImplementedError

    def get_available_data_asset_names(self, batch_kwargs_generator_names=None):
        """
        Returns a dictionary of data_asset_names that the specified batch kwarg
        generator can provide. Note that some batch kwargs generators may not be
        capable of describing specific named data assets, and some (such as
        filesystem glob batch kwargs generators) require the user to configure
        data asset names.

        Args:
            batch_kwargs_generator_names: the BatchKwargGenerator for which to get available data asset names.

        Returns:
            dictionary consisting of sets of generator assets available for the specified generators:
            ::

                {
                  generator_name: {
                    names: [ (data_asset_1, data_asset_1_type), (data_asset_2, data_asset_2_type) ... ]
                  }
                  ...
                }

        """
        available_data_asset_names = {}
        if batch_kwargs_generator_names is None:
            batch_kwargs_generator_names = [generator["name"] for generator in self.list_batch_kwargs_generators()]
        elif isinstance(batch_kwargs_generator_names, str):
            batch_kwargs_generator_names = [batch_kwargs_generator_names]

        for generator_name in batch_kwargs_generator_names:
            generator = self.get_batch_kwargs_generator(generator_name)
            available_data_asset_names[generator_name] = generator.get_available_data_asset_names()
        return available_data_asset_names

    def build_batch_kwargs(self, batch_kwargs_generator, name=None, partition_id=None, **kwargs):
        generator_obj = self.get_batch_kwargs_generator(batch_kwargs_generator)
        if partition_id is not None:
            kwargs["partition_id"] = partition_id
        return generator_obj.build_batch_kwargs(name=name, **kwargs)
