from io import StringIO
from copy import deepcopy
from botocore.exceptions import ClientError
import boto3
try:
    import cPickel as pickle
except ImportError:
    import pickle

from ds_foundation.handlers.abstract_handlers import AbstractSourceHandler, ConnectorContract, AbstractPersistHandler
from ds_connectors.parsers.dsv import DelimitedParser

__author__ = 'Darryl Oatridge, Neil Pasricha'


class AwsS3SourceHandler(AbstractSourceHandler):
    """ An AWS S3 source handler

        Delimitor Separated Value (dsv):
            encoding: the source encoding format. defualt utf-8
            file_stream: a String IO file stream to parse
            fieldnames: a set of field names for the dictionary if no header
            restkey: If row has more fields than fieldnames, remaining data is put in a list under restkey
            restval: If non-blank row has fewer fields than fieldnames, values are filled-in with restval
            dialect: Dialect class grouping formatting parameters
            delimiter: A one-character string used to separate fields. It defaults to ','.
            quotechar: A one-character string used to quote fields containing special characters, such as the
                        delimiter or quotechar, or which contain new-line characters. It defaults to '"'.
            escapechar: A one-character string used by the writer to escape the delimiter if quoting is set to
                        QUOTE_NONE and the quotechar if doublequote is False.
            doublequote: Controls how instances of quotechar appearing inside a field should themselves be quoted.
                        When True, the character is doubled. Default
                        When False, the escapechar is used as a prefix to the quotechar
            skipinitialspace: When True, whitespace immediately following the delimiter is ignored. Default False
            lineterminator: The string used to terminate lines produced by the writer.
            quoting: Controls when quotes should be generated by the writer and recognised by the reader
    """

    def __init__(self, connector_contract: ConnectorContract):
        """ initialise the Hander passing the connector_contract dictionary """
        super().__init__(connector_contract)
        self._modified = 0

    def supported_types(self) -> list:
        """ The source types supported with this module"""
        return ['csv', 'dsv', 'pickle']

    def load_canonical(self) -> dict:
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        connector_type = self.connector_contract.connector_type
        bucket = self.connector_contract.location
        _kwargs = self.connector_contract.kwargs
        if connector_type.lower() not in self.supported_types():
            raise ValueError(f"The source type '{connector_type}' is not supported. see supported_types()")
        s3_client = boto3.client('s3')
        try:
            s3_object = s3_client.get_object(Bucket=bucket, Key=resource)
        except ClientError as e:
            code = e.response["Error"]["Code"]
            raise ConnectionError(f"Failed to retrieve the object from S3 client with error code '{code}")
        _kwargs = deepcopy(self.connector_contract.kwargs)
        if connector_type.lower() in ['dsv', 'csv']:
            encoding = 'utf-8' if _kwargs.get('encoding') is None else _kwargs.pop('encoding')
            resource_body = s3_object['Body'].read().decode(encoding)
            return DelimitedParser.read_dsv(StringIO(resource_body), **_kwargs)
        if connector_type.lower() in ['pickle']:
            return pickle.load(s3_object['Body'].read())
        raise LookupError(f'The connector type {connector_type} is not currently supported')

    def get_modified(self) -> [int, float, str]:
        """ returns if the file has been modified"""
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        bucket = self.connector_contract.location
        s3_client = boto3.client('s3')
        try:
            s3_object = s3_client.get_object(Bucket=bucket, Key=resource)
        except ClientError as e:
            code = e.response["Error"]["Code"]
            raise ConnectionError(f"Failed to retrieve the object from S3 client with error code '{code}")
        return s3_object.get('LastModified')

    @staticmethod
    def list_bucket(bucket: str, starts_with: str=None, ends_with: str=None) -> list:
        """ Returns True is the file exists """
        s3_client = boto3.client('s3')
        response = s3_client.list_objects_v2(Bucket=bucket)
        rtn_list = []
        for obj in response.get('Contents', []):
            key = str(obj.get('Key'))
            if isinstance(starts_with, str) and not key.startswith(starts_with):
                continue
            if isinstance(ends_with, str) and not key.endswith(ends_with):
                continue
            rtn_list.append(key)
        return rtn_list


class AwsS3PersistHandler(AbstractPersistHandler):

    def __init__(self, connector_contract: ConnectorContract):
        """ initialise the Hander passing the connector_contract dictionary """
        super().__init__(connector_contract)
        self._modified = 0

    def supported_types(self) -> list:
        """ The source types supported with this module"""
        return ['pickle']

    def get_modified(self) -> [int, float, str]:
        """ returns if the file has been modified"""
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        bucket = self.connector_contract.location
        s3_client = boto3.client('s3')
        try:
            s3_object = s3_client.get_object(Bucket=bucket, Key=resource)
        except ClientError as e:
            code = e.response["Error"]["Code"]
            raise ConnectionError(f"Failed to retrieve the object from S3 client with error code '{code}")
        return s3_object.get('LastModified')

    def exists(self) -> bool:
        """ Returns True is the file exists """
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        bucket = self.connector_contract.location
        s3_client = boto3.client('s3')
        response = s3_client.list_objects_v2(Bucket=bucket)
        for obj in response.get('Contents', []):
            if obj['Key'] == resource:
                return True
        return False

    def load_canonical(self) -> dict:
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        connector_type = self.connector_contract.connector_type
        bucket = self.connector_contract.location
        _kwargs = self.connector_contract.kwargs
        if connector_type.lower() not in self.supported_types():
            raise ValueError(f"The source type '{connector_type}' is not supported. see supported_types()")
        s3_client = boto3.client('s3')
        try:
            s3_object = s3_client.get_object(Bucket=bucket, Key=resource)
        except ClientError as e:
            code = e.response["Error"]["Code"]
            raise ConnectionError(f"Failed to retrieve the object from S3 client with error code '{code}")
        _kwargs = deepcopy(self.connector_contract.kwargs)
        rtn_dict = None
        if connector_type.lower() in ['pickle']:
            return pickle.loads(s3_object['Body'].read())
        raise LookupError(f'The connector type {connector_type} is not currently supported')

    def persist_canonical(self, canonical: dict) -> bool:
        """ persists either the canonical dataset or the YAML contract dictionary"""
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        connector_type = self.connector_contract.connector_type
        bucket = self.connector_contract.location
        _kwargs = self.connector_contract.kwargs
        if connector_type.lower() not in self.supported_types():
            raise ValueError(f"The source type '{connector_type}' is not supported. see supported_types()")
        s3_resource = boto3.resource('s3')
        # create the bucket with the region from the Session
        _kwargs = self.connector_contract.kwargs
        upload_args = {'ACL': 'public-read'}
        if _kwargs.get('ServerSideEncryption') is not None:
            upload_args['ServerSideEncryption'] = _kwargs.get('ServerSideEncryption')
        if _kwargs.get('StorageClass') is not None:
            upload_args['StorageClass'] = _kwargs.get('StorageClass')
        # pickle
        if connector_type.lower() == 'pickle':
            if isinstance(_kwargs, dict) and _kwargs.get('protocol') is not None:
                protocol = _kwargs.get('protocol')
            else:
                protocol = pickle.HIGHEST_PROTOCOL
            pickle_byte_obj = pickle.dumps(canonical, protocol=protocol)
            s3_resource.Object(bucket, resource).put(Body=pickle_byte_obj)
            return True
        # not found
        raise ValueError(f"PandasPersistHandler only supports 'pickle' source type,"
                         f" '{connector_type}' found in Source Contract source-type")

    def remove_canonical(self) -> bool:
        if not isinstance(self.connector_contract, ConnectorContract):
            raise ValueError("The S3 Source Connector Contract has not been set correctly")
        resource = self.connector_contract.resource
        bucket = self.connector_contract.location
        s3_client = boto3.client('s3')
        response = s3_client.response = s3_client.delete_object(Bucket=bucket, Key=resource)
        if response.get('RequestCharged') is None:
            return False
        return True

    def backup_canonical(self, max_backups=None):
        """ creates a backup of the current source contract resource"""
        if not isinstance(self.connector_contract, ConnectorContract):
            return
        max_backups = max_backups if isinstance(max_backups, int) else 10
        resource = self.connector_contract.resource
        location = self.connector_contract.location
        # _filepath = os.path.join(location, resource)
        # # Check existence of previous versions
        # name, _, ext = _filepath.rpartition('.')
        # for index in range(max_backups):
        #     backup = '%s_%2.2d.%s' % (name, index, ext)
        #     if index > 0:
        #         # No need to backup if file and last version
        #         # are identical
        #         old_backup = '%s_%2.2d.%s' % (name, index - 1, ext)
        #         if not os.path.exists(old_backup):
        #             break
        #         abspath = os.path.abspath(old_backup)
        #
        #         try:
        #             if os.path.isfile(abspath) and filecmp.cmp(abspath, _filepath, shallow=False):
        #                 continue
        #         except OSError:
        #             pass
        #     try:
        #         if not os.path.exists(backup):
        #             shutil.copy(_filepath, backup)
        #     except (OSError, IOError):
        #         pass
        # return
