import logging
import copy
from hashlib import md5
import datetime
import time

import pandas as pd
from six import string_types
from great_expectations.types import RequiredKeysDotDict, AllowedKeysDotDict, ClassConfig
from great_expectations.datasource.types.reader_methods import ReaderMethods
from great_expectations.data_context.types.base_resource_identifiers import (
    OrderedDataContextKey,
    DataContextKey,
)
# from great_expectations.exceptions import GreatExpectationsError

try:
    import pyspark
except ImportError:
    pyspark = None

logger = logging.getLogger(__name__)


class BatchFingerprint(OrderedDataContextKey):
    _allowed_keys = OrderedDataContextKey._allowed_keys | {
        "partition_id",
        "fingerprint"
    }
    _required_keys = OrderedDataContextKey._required_keys | {
        "partition_id",
        "fingerprint"
    }
    _key_types = copy.copy(OrderedDataContextKey._key_types)
    _key_types.update({
        "partition_id": string_types,
        "fingerprint": string_types
    })
    _key_order = copy.copy(OrderedDataContextKey._key_order)
    _key_order.extend(["partition_id", "fingerprint"])


class BatchKwargs(RequiredKeysDotDict):
    """BatchKwargs represent information required by a datasource to fetch a batch of data.

    BatchKwargs are usually generated by BatchGenerator objects and interpreted by Datasource objects.
    """
    _required_keys = set()
    _partition_id_key = "partition_id"  # a partition id can be used as shorthand to access a batch of data

    # _batch_fingerprint_ignored_keys makes it possible to define keys which, if present, are ignored for purposes
    # of determining the unique batch id, such that batches differing only in the value in these keys are given
    # the same id
    _batch_fingerprint_ignored_keys = {
        "data_asset_type"
    }
    _key_types = {
        "data_asset_type": ClassConfig
    }

    @property
    def batch_fingerprint(self):
        partition_id = self.get(self._partition_id_key, None)
        # We do not allow a "None" partition_id, even if it's explicitly present as such in batch_kwargs
        if partition_id is None:
            partition_id = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%S.%fZ")
        id_keys = (set(self.keys()) - set(self._batch_fingerprint_ignored_keys)) - {self._partition_id_key}
        if len(id_keys) == 1:
            key = list(id_keys)[0]
            hash_ = key + ":" + str(self[key])
        else:
            hash_dict = {k: self[k] for k in id_keys}
            hash_ = md5(str(sorted(hash_dict.items())).encode("utf-8")).hexdigest()

        return BatchFingerprint(partition_id=partition_id, fingerprint=hash_)

    @classmethod
    def build_batch_fingerprint(cls, dict_):
        try:
            return BatchKwargs(dict_).batch_fingerprint
        except (KeyError, TypeError):
            logger.warning("Unable to build BatchKwargs from provided dictionary.")
            return None


class BatchId(BatchKwargs):
    """A BatchId is a special type of BatchKwargs (so that it has a batch_fingerprint) but it generally does
    NOT require specific keys and instead captures information about the OUTPUT of a datasource's fetch
    process, such as the timestamp at which a query was executed."""
    _required_keys = BatchKwargs._required_keys | {
        "timestamp"
    }
    _key_types = copy.copy(BatchKwargs._key_types)
    _key_types.update({
        "timestamp": float
    })


class PandasDatasourceBatchKwargs(BatchKwargs):
    """This is an abstract class and should not be instantiated. It's relevant for testing whether
    a subclass is allowed
    """
    pass


class SparkDFDatasourceBatchKwargs(BatchKwargs):
    """This is an abstract class and should not be instantiated. It's relevant for testing whether
    a subclass is allowed
    """
    pass


class SqlAlchemyDatasourceBatchKwargs(BatchKwargs):
    """This is an abstract class and should not be instantiated. It's relevant for testing whether
    a subclass is allowed
    """
    pass


class PathBatchKwargs(PandasDatasourceBatchKwargs, SparkDFDatasourceBatchKwargs):
    """PathBatchKwargs represents kwargs suitable for reading a file from a given path."""
    _required_keys = {
        "path"
    }
    # NOTE: JPC - 20190821: Eventually, we will probably want to have some logic that decides to use, say,
    # an md5 hash of a file instead of a path to decide when it's the same, or to differentiate paths
    # from s3 from paths on a local filesystem
    _key_types = {
        "path": string_types,
        "reader_method": ReaderMethods
    }


class S3BatchKwargs(PandasDatasourceBatchKwargs, SparkDFDatasourceBatchKwargs):
    """PathBatchKwargs represents kwargs suitable for reading a file from a given path."""
    _required_keys = {
        "s3"
    }
    # NOTE: JPC - 20190821: Eventually, we will probably want to have some logic that decides to use, say,
    # an md5 hash of a file instead of a path to decide when it's the same, or to differentiate paths
    # from s3 from paths on a local filesystem
    _key_types = {
        "s3": string_types,
        "reader_method": ReaderMethods
    }


class InMemoryBatchKwargs(PandasDatasourceBatchKwargs, SparkDFDatasourceBatchKwargs):
    _required_keys = {
        "dataset"
    }


class PandasDatasourceMemoryBatchKwargs(InMemoryBatchKwargs):
    _key_types = {
        "dataset": pd.DataFrame
    }


class SparkDFDatasourceMemoryBatchKwargs(InMemoryBatchKwargs):
    try:
        _key_types = {
            "dataset": pyspark.sql.DataFrame
        }
    except AttributeError:
        _key_types = {
            "dataset": None  # If we were unable to import pyspark, these are invalid
        }


class SqlAlchemyDatasourceTableBatchKwargs(SqlAlchemyDatasourceBatchKwargs):
    _required_keys = {
        "table"
    }
    _key_types = {
        "table": string_types
    }


class SqlAlchemyDatasourceQueryBatchKwargs(SqlAlchemyDatasourceBatchKwargs):
    _required_keys = {
        "query"
    }
    _key_types = {
        "query": string_types
    }


class SparkDFDatasourceQueryBatchKwargs(SparkDFDatasourceBatchKwargs):
    _required_keys = {
        "query"
    }
    _key_types = {
        "query": string_types
    }
