"""S3 connector."""

import hashlib

import aioboto3

import config

import src.exceptions


class S3Connector:
    """S3 operations."""

    session = None
    download_chunk_size = None
    signed_url_expires = None

    def __init__(self):
        """Initialize AWS session."""
        self.download_chunk_size = 1024 * config.DOWNLOAD_CHUNK_SIZE_MB
        self.signed_url_expires = config.S3_URL_EXPIRES
        self.session = aioboto3.Session()

    async def list_objects(self, bucket, prefix=''):
        """Get object keys in bucket.

        Args:
            bucket (str): location to search
            prefix (str): sub location to search

        Returns:
            (list): all keys in bucket
        """
        keys = list()
        async with self.session.client('s3') as client:
            paginator = client.get_paginator('list_objects')
            async for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
                for content in page.get('Contents', []):
                    key = content['Key']
                    if not key.endswith('/'):
                        keys.append(key)
        return keys

    async def upload_object(self, data, bucket, path):
        """Upload local file or data to s3.

        Args:
            data (TextIOWrapper|bytes): file handler or data to upload
            bucket (str): s3 bucket destination
            path (str): s3 path destination

        Returns:
            None
        """
        async with self.session.client('s3') as client:
            await client.put_object(
                Body=data,
                Bucket=bucket,
                Key=path
            )

    async def copy_object(self, from_bucket, from_path, to_bucket, to_path):
        """Copy s3 to s3 via multipart upload.

        Args:
            from_bucket (str): source bucket
            from_path (str): source path
            to_bucket (str): destination bucket
            to_path (str): destination path

        Returns:
            None
        """
        async with self.session.client('s3') as client:
            part_size = 500 * 1024 * 1024  # 500 MB
            part_number = 1
            parts = []
            source_object_size = (await client.head_object(
                Bucket=from_bucket, Key=from_path))['ContentLength']

            # Initiate multipart upload
            multipart_upload = await client.create_multipart_upload(
                Bucket=to_bucket,
                Key=to_path
            )

            try:
                for start_byte in range(0, source_object_size, part_size):
                    end_byte = min(start_byte + part_size - 1, source_object_size - 1)  # noqa:E501

                    copy_part_response = await client.upload_part_copy(
                        Bucket=to_bucket,
                        Key=to_path,
                        CopySource={'Bucket': from_bucket, 'Key': from_path},
                        CopySourceRange=f'bytes={start_byte}-{end_byte}',
                        PartNumber=part_number,
                        UploadId=multipart_upload['UploadId'],
                    )

                    parts.append({
                        'PartNumber': part_number,
                        'ETag': copy_part_response['CopyPartResult']['ETag']
                    })

                    part_number += 1

                # Complete multipart upload
                await client.complete_multipart_upload(
                    Bucket=to_bucket,
                    Key=to_path,
                    UploadId=multipart_upload['UploadId'],
                    MultipartUpload={'Parts': parts}
                )

            except Exception as upload_exception:
                try:
                    await client.abort_multipart_upload(
                        Bucket=to_bucket,
                        Key=to_path,
                        UploadId=multipart_upload['UploadId'],
                    )
                except Exception as abort_upload_exception:
                    raise src.exceptions.CopyAudioFileToOutputFailedError(
                        f'Failed to upload {from_bucket}/{from_path} to {to_bucket}/{to_path}. '  # noqa:E501
                        f'Errors: {repr(upload_exception)}, {repr(abort_upload_exception)}'  # noqa:E501
                    )

                raise src.exceptions.CopyAudioFileToOutputFailedError(
                    f'Failed to upload {from_bucket}/{from_path} to {to_bucket}/{to_path}. '  # noqa:E501
                    f'Error: {repr(upload_exception)}'
                )

    async def analyze_object(self, bucket, path):
        """Calculate MD5 of object.

        Args:
            bucket (str): bucket containing object
            path (str): location of object in bucket

        Returns:
            (tuple):
                0 (str): bucket location of object
                1 (str): key location of object
                2 (str): md5 of object body
                3 (int): size of object body
        """
        md5_hash = hashlib.md5()
        async with self.session.client('s3') as client:
            file_data = await client.get_object(Bucket=bucket, Key=path)
            async for chunk in file_data['Body'].iter_chunks(chunk_size=self.download_chunk_size):  # noqa:E501
                md5_hash.update(chunk)
        return (
            bucket,
            path,
            md5_hash.hexdigest(),
            file_data['ContentLength']
        )

    async def get_signed_url(self, bucket, path):
        """Generate expiring s3 url for get_object.

        Args:
            bucket (str): bucket containing object
            path (str): location of object in bucket

        Returns:
            (tuple):
                0 (str): bucket location of object
                1 (str): key location of object
                2 (str): signed url to fetch object at
        """
        async with self.session.client('s3') as client:
            url = await client.generate_presigned_url(
                ClientMethod='get_object',
                ExpiresIn=self.signed_url_expires,
                Params={
                    'Bucket': bucket,
                    'Key': path
                }
            )
        return (
            bucket,
            path,
            url
        )

    async def download_fileobj(self, bucket, path, output_filename):
        """Download s3 object to file on disk.

        Args:
            bucket (str): bucket containing object
            path (str): location of object in bucket
            output_filename (str): location of output file on disk

        Returns:
            None
        """
        async with self.session.client('s3') as client:
            with open(output_filename, 'wb') as data:
                await client.download_fileobj(bucket, path, data)