# Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
import copy
import math
from urllib.parse import parse_qsl

from botocore.exceptions import ClientError

from s3transfer.exceptions import S3CopyFailedError
from s3transfer.tasks import (
    CompleteMultipartUploadTask,
    CreateMultipartUploadTask,
    SubmissionTask,
    Task,
)
from s3transfer.utils import (
    ChunksizeAdjuster,
    calculate_range_parameter,
    get_callbacks,
    get_filtered_dict,
)


class CopySubmissionTask(SubmissionTask):
    """Task for submitting tasks to execute a copy"""

    EXTRA_ARGS_TO_HEAD_ARGS_MAPPING = {
        'CopySourceIfMatch': 'IfMatch',
        'CopySourceIfModifiedSince': 'IfModifiedSince',
        'CopySourceIfNoneMatch': 'IfNoneMatch',
        'CopySourceIfUnmodifiedSince': 'IfUnmodifiedSince',
        'CopySourceSSECustomerKey': 'SSECustomerKey',
        'CopySourceSSECustomerAlgorithm': 'SSECustomerAlgorithm',
        'CopySourceSSECustomerKeyMD5': 'SSECustomerKeyMD5',
        'RequestPayer': 'RequestPayer',
        'ExpectedBucketOwner': 'ExpectedBucketOwner',
    }

    UPLOAD_PART_COPY_ARGS = [
        'CopySourceIfMatch',
        'CopySourceIfModifiedSince',
        'CopySourceIfNoneMatch',
        'CopySourceIfUnmodifiedSince',
        'CopySourceSSECustomerKey',
        'CopySourceSSECustomerAlgorithm',
        'CopySourceSSECustomerKeyMD5',
        'SSECustomerKey',
        'SSECustomerAlgorithm',
        'SSECustomerKeyMD5',
        'RequestPayer',
        'ExpectedBucketOwner',
    ]

    CREATE_MULTIPART_ARGS_BLACKLIST = [
        'CopySourceIfMatch',
        'CopySourceIfModifiedSince',
        'CopySourceIfNoneMatch',
        'CopySourceIfUnmodifiedSince',
        'CopySourceSSECustomerKey',
        'CopySourceSSECustomerAlgorithm',
        'CopySourceSSECustomerKeyMD5',
        'MetadataDirective',
        'TaggingDirective',
        'AnnotationDirective',
        'Tagging',
    ]

    # Metadata fields to preserve for multipart copies.
    PRESERVED_METADATA_FIELDS = [
        'CacheControl',
        'ContentDisposition',
        'ContentEncoding',
        'ContentLanguage',
        'ContentType',
        'Expires',
        'Metadata',
    ]

    COMPLETE_MULTIPART_ARGS = [
        'SSECustomerKey',
        'SSECustomerAlgorithm',
        'SSECustomerKeyMD5',
        'RequestPayer',
        'ExpectedBucketOwner',
    ]

    GET_OBJECT_TAGGING_ARGS = ['RequestPayer', 'ExpectedBucketOwner']
    PUT_OBJECT_TAGGING_ARGS = [
        'RequestPayer',
        'ExpectedBucketOwner',
        'ChecksumAlgorithm',
    ]
    LIST_OBJECT_ANNOTATIONS_ARGS = ['RequestPayer', 'ExpectedBucketOwner']
    GET_OBJECT_ANNOTATION_ARGS = ['RequestPayer', 'ExpectedBucketOwner']
    PUT_OBJECT_ANNOTATION_ARGS = [
        'RequestPayer',
        'ExpectedBucketOwner',
        'ChecksumAlgorithm',
    ]

    def _submit(
        self, client, config, osutil, request_executor, transfer_future
    ):
        """
        :param client: The client associated with the transfer manager

        :type config: s3transfer.manager.TransferConfig
        :param config: The transfer config associated with the transfer
            manager

        :type osutil: s3transfer.utils.OSUtil
        :param osutil: The os utility associated to the transfer manager

        :type request_executor: s3transfer.futures.BoundedExecutor
        :param request_executor: The request executor associated with the
            transfer manager

        :type transfer_future: s3transfer.futures.TransferFuture
        :param transfer_future: The transfer future associated with the
            transfer request that tasks are being submitted for
        """
        preserved_metadata = {}
        source_version_id = None
        call_args = transfer_future.meta.call_args
        if (
            transfer_future.meta.size is None
            or transfer_future.meta.etag is None
        ):
            # If a size was not provided figure out the size for the
            # user. Note that we will only use the client provided to
            # the TransferManager. If the object is outside of the region
            # of the client, they may have to provide the file size themselves
            # with a completely new client.
            head_object_request = (
                self._get_head_object_request_from_copy_source(
                    call_args.copy_source
                )
            )
            extra_args = call_args.extra_args

            # Map any values that may be used in the head object that is
            # used in the copy object
            for param, value in extra_args.items():
                if param in self.EXTRA_ARGS_TO_HEAD_ARGS_MAPPING:
                    head_object_request[
                        self.EXTRA_ARGS_TO_HEAD_ARGS_MAPPING[param]
                    ] = value

            response = call_args.source_client.head_object(
                **head_object_request
            )
            transfer_future.meta.provide_transfer_size(
                response['ContentLength']
            )
            # Provide an etag to ensure a stored object is not modified
            # during a multipart copy.
            transfer_future.meta.provide_object_etag(response.get('ETag'))
            preserved_metadata = self._extract_preserved_metadata(response)
            # Pin the source version so all subsequent reads (tags, annotations)
            # are consistent with the object from the head call
            source_version_id = response.get('VersionId')

        # If it is greater than threshold do a multipart copy, otherwise
        # do a regular copy object.
        if transfer_future.meta.size < config.multipart_threshold:
            self._submit_copy_request(
                client, config, osutil, request_executor, transfer_future
            )
        else:
            self._submit_multipart_request(
                client,
                config,
                osutil,
                request_executor,
                transfer_future,
                preserved_metadata,
                source_version_id=source_version_id,
            )

    def _submit_copy_request(
        self, client, config, osutil, request_executor, transfer_future
    ):
        call_args = transfer_future.meta.call_args

        # Get the needed progress callbacks for the task
        progress_callbacks = get_callbacks(transfer_future, 'progress')

        # Submit the request of a single copy.
        self._transfer_coordinator.submit(
            request_executor,
            CopyObjectTask(
                transfer_coordinator=self._transfer_coordinator,
                main_kwargs={
                    'client': client,
                    'copy_source': call_args.copy_source,
                    'bucket': call_args.bucket,
                    'key': call_args.key,
                    'extra_args': call_args.extra_args,
                    'callbacks': progress_callbacks,
                    'size': transfer_future.meta.size,
                },
                is_final=True,
            ),
        )

    def _submit_multipart_request(
        self,
        client,
        config,
        osutil,
        request_executor,
        transfer_future,
        preserved_metadata=None,
        source_version_id=None,
    ):
        call_args = transfer_future.meta.call_args
        merged_extra_args = self._apply_preserved_metadata(
            call_args.extra_args, preserved_metadata or {}
        )

        # Submit the request to create a multipart upload and make sure it
        # does not include any of the arguments used for copy part.
        create_multipart_extra_args = {}
        for param, val in merged_extra_args.items():
            if param not in self.CREATE_MULTIPART_ARGS_BLACKLIST:
                create_multipart_extra_args[param] = val

        create_multipart_future = self._transfer_coordinator.submit(
            request_executor,
            CreateMultipartUploadTask(
                transfer_coordinator=self._transfer_coordinator,
                main_kwargs={
                    'client': client,
                    'bucket': call_args.bucket,
                    'key': call_args.key,
                    'extra_args': create_multipart_extra_args,
                },
            ),
        )

        # Determine how many parts are needed based on filesize and
        # desired chunksize.
        part_size = config.multipart_chunksize
        adjuster = ChunksizeAdjuster()
        part_size = adjuster.adjust_chunksize(
            part_size, transfer_future.meta.size
        )
        num_parts = int(
            math.ceil(transfer_future.meta.size / float(part_size))
        )

        # Submit requests to upload the parts of the file.
        part_futures = []
        progress_callbacks = get_callbacks(transfer_future, 'progress')

        for part_number in range(1, num_parts + 1):
            extra_part_args = self._extra_upload_part_args(
                call_args.extra_args
            )
            # The part number for upload part starts at 1 while the
            # range parameter starts at zero, so just subtract 1 off of
            # the part number
            extra_part_args['CopySourceRange'] = calculate_range_parameter(
                part_size,
                part_number - 1,
                num_parts,
                transfer_future.meta.size,
            )
            if transfer_future.meta.etag is not None:
                extra_part_args['CopySourceIfMatch'] = (
                    transfer_future.meta.etag
                )
            # Get the size of the part copy as well for the progress
            # callbacks.
            size = self._get_transfer_size(
                part_size,
                part_number - 1,
                num_parts,
                transfer_future.meta.size,
            )
            # Get the checksum algorithm of the multipart request.
            checksum_algorithm = call_args.extra_args.get("ChecksumAlgorithm")
            part_futures.append(
                self._transfer_coordinator.submit(
                    request_executor,
                    CopyPartTask(
                        transfer_coordinator=self._transfer_coordinator,
                        main_kwargs={
                            'client': client,
                            'copy_source': call_args.copy_source,
                            'bucket': call_args.bucket,
                            'key': call_args.key,
                            'part_number': part_number,
                            'extra_args': extra_part_args,
                            'callbacks': progress_callbacks,
                            'size': size,
                            'checksum_algorithm': checksum_algorithm,
                        },
                        pending_main_kwargs={
                            'upload_id': create_multipart_future
                        },
                    ),
                )
            )

        complete_multipart_extra_args = self._extra_complete_multipart_args(
            call_args.extra_args
        )

        # Submit the request to complete the multipart upload.
        self._transfer_coordinator.submit(
            request_executor,
            CopyCompleteMultipartUploadTask(
                transfer_coordinator=self._transfer_coordinator,
                main_kwargs={
                    'client': client,
                    'bucket': call_args.bucket,
                    'key': call_args.key,
                    'extra_args': complete_multipart_extra_args,
                    'call_args': call_args,
                    'source_version_id': source_version_id,
                },
                pending_main_kwargs={
                    'upload_id': create_multipart_future,
                    'parts': part_futures,
                },
                is_final=True,
            ),
        )

    def _extract_preserved_metadata(self, head_object_response):
        preserved = {}
        for field in self.PRESERVED_METADATA_FIELDS:
            if field in head_object_response:
                preserved[field] = head_object_response[field]
        return preserved

    def _apply_preserved_metadata(self, extra_args, preserved_metadata):
        # MPU has no native MetadataDirective, handle metadata manually.  REPLACE
        # means we copy whatever the user provided, anything else means we drop
        # what the user supplied
        if extra_args.get('MetadataDirective') == 'REPLACE':
            return extra_args
        result = {
            k: v for k, v in extra_args.items()
            if k not in self.PRESERVED_METADATA_FIELDS
        }
        result.update(preserved_metadata)
        return result

    def _get_head_object_request_from_copy_source(self, copy_source):
        if isinstance(copy_source, dict):
            return copy.copy(copy_source)
        else:
            raise TypeError(
                'Expecting dictionary formatted: '
                '{"Bucket": bucket_name, "Key": key} '
                f'but got {copy_source} or type {type(copy_source)}.'
            )

    def _extra_upload_part_args(self, extra_args):
        # Only the args in COPY_PART_ARGS actually need to be passed
        # onto the upload_part_copy calls.
        return get_filtered_dict(extra_args, self.UPLOAD_PART_COPY_ARGS)

    def _extra_complete_multipart_args(self, extra_args):
        return get_filtered_dict(extra_args, self.COMPLETE_MULTIPART_ARGS)

    def _get_transfer_size(
        self, part_size, part_index, num_parts, total_transfer_size
    ):
        if part_index == num_parts - 1:
            # The last part may be different in size then the rest of the
            # parts.
            return total_transfer_size - (part_index * part_size)
        return part_size


class CopyCompleteMultipartUploadTask(CompleteMultipartUploadTask):
    """CompleteMultipartUpload variant that also applies tags and annotations.

    After the destination object is finalized, copies/applies tags and
    annotations inline. Errors during apply propagate as task failures.
    """

    def _main(
        self,
        client,
        bucket,
        key,
        upload_id,
        parts,
        extra_args,
        call_args,
        source_version_id,
    ):
        response = client.complete_multipart_upload(
            Bucket=bucket,
            Key=key,
            UploadId=upload_id,
            MultipartUpload={'Parts': parts},
            **extra_args,
        )
        dest_etag = response.get('ETag')
        dest_version_id = response.get('VersionId')
        self._apply_tags(client, call_args, source_version_id, dest_version_id)
        self._apply_annotations(
            client, call_args, source_version_id, dest_version_id, dest_etag
        )

    def _apply_tags(self, client, call_args, source_version_id, dest_version_id):
        extra_args = call_args.extra_args
        directive = extra_args.get('TaggingDirective')
        if directive not in ('COPY', 'REPLACE'):
            return
        if directive == 'COPY':
            src_kwargs = {
                'Bucket': call_args.copy_source['Bucket'],
                'Key': call_args.copy_source['Key'],
                **get_filtered_dict(
                    extra_args, CopySubmissionTask.GET_OBJECT_TAGGING_ARGS
                ),
            }
            if source_version_id:
                src_kwargs['VersionId'] = source_version_id
            tag_set = call_args.source_client.get_object_tagging(
                **src_kwargs
            ).get('TagSet', [])
        else:  # REPLACE
            tag_set = [
                {'Key': k, 'Value': v}
                for k, v in parse_qsl(
                    extra_args.get('Tagging', ''),
                    keep_blank_values=True,
                )
            ]
        if not tag_set:
            return
        put_kwargs = {
            'Bucket': call_args.bucket,
            'Key': call_args.key,
            'Tagging': {'TagSet': tag_set},
            **get_filtered_dict(
                extra_args, CopySubmissionTask.PUT_OBJECT_TAGGING_ARGS
            ),
        }
        if dest_version_id:
            put_kwargs['VersionId'] = dest_version_id
        client.put_object_tagging(**put_kwargs)

    def _apply_annotations(
        self,
        client,
        call_args,
        source_version_id,
        dest_version_id,
        dest_etag,
    ):
        # We copy annotations only if COPY is explicitly set by the user.
        extra_args = call_args.extra_args
        if extra_args.get('AnnotationDirective') != 'COPY':
            return
        src_base = {
            'Bucket': call_args.copy_source['Bucket'],
            'Key': call_args.copy_source['Key'],
        }
        if source_version_id:
            src_base['VersionId'] = source_version_id
        list_kwargs = {
            **src_base,
            **get_filtered_dict(
                extra_args, CopySubmissionTask.LIST_OBJECT_ANNOTATIONS_ARGS
            ),
        }
        get_kwargs_base = {
            **src_base,
            **get_filtered_dict(
                extra_args, CopySubmissionTask.GET_OBJECT_ANNOTATION_ARGS
            ),
        }
        put_passthrough = get_filtered_dict(
            extra_args, CopySubmissionTask.PUT_OBJECT_ANNOTATION_ARGS
        )
        list_response = call_args.source_client.list_object_annotations(
            **list_kwargs
        )
        succeeded = []
        failed = {}
        for annotation in list_response.get('Annotations', []):
            name = annotation['AnnotationName']
            payload_response = call_args.source_client.get_object_annotation(
                **get_kwargs_base,
                AnnotationName=name,
            )
            put_kwargs = {
                'Bucket': call_args.bucket,
                'Key': call_args.key,
                'AnnotationName': name,
                'AnnotationPayload': payload_response['AnnotationPayload'].read(),
                **put_passthrough,
            }
            if dest_version_id:
                put_kwargs['VersionId'] = dest_version_id
            if dest_etag:
                put_kwargs['ObjectIfMatch'] = dest_etag
            try:
                client.put_object_annotation(**put_kwargs)
                succeeded.append(name)
            except Exception as e:
                failed[name] = e
        if failed:
            raise S3CopyFailedError(
                f'Failed to copy annotations to '
                f's3://{call_args.bucket}/{call_args.key}. '
                f'Succeeded: {succeeded}. '
                f'Failed: {list(failed.keys())}. '
                f'Errors: {failed}'
            )


class CopyObjectTask(Task):
    """Task to do a nonmultipart copy"""

    def _main(
        self, client, copy_source, bucket, key, extra_args, callbacks, size
    ):
        """
        :param client: The client to use when calling PutObject
        :param copy_source: The CopySource parameter to use
        :param bucket: The name of the bucket to copy to
        :param key: The name of the key to copy to
        :param extra_args: A dictionary of any extra arguments that may be
            used in the upload.
        :param callbacks: List of callbacks to call after copy
        :param size: The size of the transfer. This value is passed into
            the callbacks

        """
        client.copy_object(
            CopySource=copy_source, Bucket=bucket, Key=key, **extra_args
        )
        for callback in callbacks:
            callback(bytes_transferred=size)


class CopyPartTask(Task):
    """Task to upload a part in a multipart copy"""

    def _main(
        self,
        client,
        copy_source,
        bucket,
        key,
        upload_id,
        part_number,
        extra_args,
        callbacks,
        size,
        checksum_algorithm=None,
    ):
        """
        :param client: The client to use when calling PutObject
        :param copy_source: The CopySource parameter to use
        :param bucket: The name of the bucket to upload to
        :param key: The name of the key to upload to
        :param upload_id: The id of the upload
        :param part_number: The number representing the part of the multipart
            upload
        :param extra_args: A dictionary of any extra arguments that may be
            used in the upload.
        :param callbacks: List of callbacks to call after copy part
        :param size: The size of the transfer. This value is passed into
            the callbacks
        :param checksum_algorithm: The algorithm that was used to create the multipart
            upload

        :rtype: dict
        :returns: A dictionary representing a part::

            {'Etag': etag_value, 'PartNumber': part_number}

            This value can be appended to a list to be used to complete
            the multipart upload. If a checksum is in the response,
            it will also be included.
        """
        try:
            response = client.upload_part_copy(
                CopySource=copy_source,
                Bucket=bucket,
                Key=key,
                UploadId=upload_id,
                PartNumber=part_number,
                **extra_args,
            )
        except ClientError as e:
            error_code = e.response.get('Error', {}).get('Code')
            src_key = copy_source['Key']
            src_bucket = copy_source['Bucket']
            if error_code == "PreconditionFailed":
                raise S3CopyFailedError(
                    f'Contents of stored object "{src_key}" '
                    f'in bucket "{src_bucket}" did not match '
                    'expected ETag.'
                )
            else:
                raise
        for callback in callbacks:
            callback(bytes_transferred=size)
        etag = response['CopyPartResult']['ETag']
        part_metadata = {'ETag': etag, 'PartNumber': part_number}
        if checksum_algorithm:
            checksum_member = f'Checksum{checksum_algorithm.upper()}'
            if checksum_member in response['CopyPartResult']:
                part_metadata[checksum_member] = response['CopyPartResult'][
                    checksum_member
                ]
        return part_metadata