# coding: utf-8
# Copyright (c) 2016, 2020, Oracle and/or its affiliates.  All rights reserved.
# This software is dual-licensed to you under the Universal Permissive License (UPL) 1.0 as shown at https://oss.oracle.com/licenses/upl or Apache License 2.0 as shown at http://www.apache.org/licenses/LICENSE-2.0. You may choose either license.

import io
import hashlib
import base64
from . import md5 as MD5
from multiprocessing.dummy import Pool
from os import stat
from ..constants import DEFAULT_PART_SIZE
from ..constants import MEBIBYTE
from .buffered_part_reader import BufferedPartReader
from ... import models
from ....exceptions import ServiceError, MultipartUploadError
from oci.exceptions import RequestException, ConnectTimeout
from oci._vendor.requests.exceptions import Timeout, ConnectionError
from oci._vendor.six.moves.queue import Queue
from threading import Semaphore
from oci._vendor import six
from oci.fips import is_fips_mode

READ_BUFFER_SIZE = 8 * 1024
DEFAULT_PARALLEL_PROCESS_COUNT = 3
DEFAULT_MAX_RETRIES = 3
SSEC_PARAM_NAMES = [
    'opc_sse_customer_algorithm',
    'opc_sse_customer_key',
    'opc_sse_customer_key_sha256'
]

# This timeout value will be used only for multipart upload operations, since we're running into SSL issues with
# sending memoryviews with a timeout.
MULTIPART_UPLOAD_TIMEOUT = None


class MultipartObjectAssembler:
    def __init__(self,
                 object_storage_client,
                 namespace_name,
                 bucket_name,
                 object_name,
                 **kwargs
                 ):
        """
        MultipartObjectAssembler provides a simplified interaction when uploading large
        objects using multi-part uploads.

        An assembler can be used to begin a new upload, or resume a previous one.

        PLEASE NOTE that the timeout for the object storage client is overwritten to `None` for all operations which
        call object storage. For this reason, the operations are NOT thread-safe, and you should provide the
        MultipartObjectAssembler class with its own Object Storage client that isn't used elsewhere.
        For more information please see `Known Issues <https://oracle-cloud-infrastructure-python-sdk.readthedocs.io/en/latest/known-issues.html>`_

        :param ObjectStorageClient object_storage_client:
            A configured Object Storage client.

        :param str namespace_name:
            The namespace containing the bucket in which to store the object.

        :param str bucket_name:
            The name of the bucket in which to store the object.

        :param str object_name:
            The name to use for the object in Object Storage, or existing object name if resuming.

        :param int part_size (optional):
            Override the default part size of 128 MiB, value is in bytes.

        :param str if_match (optional):
            The entity tag of the object to match.

        :param str if_none_match (optional):
            The entity tag of the object to avoid matching. The only valid value is ‘*’,
            which indicates that the request should fail if the object
             already exists.

        :param str content_type (optional):
            The content type of the object to upload.

        :param str content_language (optional):
            The content language of the object to upload.

        :param str content_encoding (optional):
            The content encoding of the object to upload.

        :param dict metadata (optional):
            A dictionary of string to string values to associate with the object to upload

        :param bool allow_parallel_uploads (optional):
            Whether or not this MultipartObjectAssembler supports uploading parts in parallel. Defaults to True.

        :param int parallel_process_count (optional):
            The number of worker processes to use in parallel for performing a multipart upload. Default is 3.

        :param str opc_sse_customer_algorithm (optional):
            The encryption algorithm to use with the customer provided encryption key.

        :param str opc_sse_customer_key (optional):
            The base64-encoded 256-bit encryption key to use to encrypt the objects uploaded by this MultipartObjectAssembler.

        :param str opc_sse_customer_key_sha256 (optional):
            The base64-encoded SHA256 hash of the encryption key.
        """
        self.object_storage_client = object_storage_client

        self.part_size = DEFAULT_PART_SIZE
        if 'part_size' in kwargs:
            self.part_size = kwargs['part_size']
        self.if_match = None
        if 'if_match' in kwargs:
            self.if_match = kwargs['if_match']
        self.if_none_match = None
        if 'if_none_match' in kwargs:
            self.if_none_match = kwargs['if_none_match']
        self.content_type = None
        if 'content_type' in kwargs:
            self.content_type = kwargs['content_type']
        self.content_language = None
        if 'content_language' in kwargs:
            self.content_language = kwargs['content_language']
        self.content_encoding = None
        if 'content_encoding' in kwargs:
            self.content_encoding = kwargs['content_encoding']
        self.metadata = None
        if 'metadata' in kwargs:
            self.metadata = kwargs['metadata']

        self.parallel_process_count = DEFAULT_PARALLEL_PROCESS_COUNT
        if 'parallel_process_count' in kwargs:
            if kwargs['parallel_process_count'] == 0:
                raise ValueError('parallel_process_count must be at least one.')

            self.parallel_process_count = kwargs['parallel_process_count']

        if 'allow_parallel_uploads' in kwargs and kwargs['allow_parallel_uploads'] is False:
            # if parallel uploads are disabled, use only a single process
            self.parallel_process_count = 1

        self.max_retries = DEFAULT_MAX_RETRIES
        self.manifest = {"uploadId": None,
                         "namespace": namespace_name,
                         "bucketName": bucket_name,
                         "objectName": object_name,
                         "parts": []}
        # Copy SSE-C parameters (if any)
        self.ssec_params = {}
        for param_name in SSEC_PARAM_NAMES:
            if param_name in kwargs:
                self.ssec_params[param_name] = kwargs[param_name]

        # Store the original timeout value
        self._object_storage_client_timeout = self.object_storage_client.base_client.timeout

    @staticmethod
    def calculate_md5(file_path, offset, chunk):
        """
        Calculate the base64 encoded MD5 hash for a part of a file.

        :param str file_path: Path to the file
        :param int offset: Offset where the part starts in the file
        :param int chunk: Number of bytes in the part
        :return: Base64 encoded MD5 hash
        :rtype: str
        """

        # Determine if we can use the hashlib version of md5 or the bundled
        # version of md5
        if is_fips_mode():
            md5 = MD5.md5
        else:
            md5 = hashlib.md5

        m = md5()
        with io.open(file_path, mode='rb') as f:
            bpr = BufferedPartReader(f, offset, chunk)
            while True:
                part = bpr.read(READ_BUFFER_SIZE)
                if part == b'':
                    break
                m.update(part)
        return base64.b64encode(m.digest()).decode("utf-8")

    @staticmethod
    def _is_exception_retryable(e):
        """
        Determines if the service should attempt to retry an operation based
         on the type of exception

        retry if
           timeout
           Connection error
           unknown client exception: status == -1
           server exception: status >= 500
           Potential edge case: status == 409

        :param e: Exception
        :return: Boolean
        """
        retryable = False
        if isinstance(e, Timeout):
            retryable = True
        elif isinstance(e, ConnectionError):
            retryable = True
        elif isinstance(e, RequestException):
            retryable = True
        elif isinstance(e, ConnectTimeout):
            retryable = True
        elif isinstance(e, ServiceError):
            if e.status >= 500 or e.status == -1 or (e.status == 409 and e.code == "ConcurrentObjectUpdate") or (e.status == 429):
                retryable = True

        return retryable

    def _upload_part_call(self, object_storage_client, **kwargs):
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()
            with io.open(kwargs["part_file_path"], mode='rb') as file_object:
                bpr = BufferedPartReader(file_object, kwargs["offset"], kwargs["size"])

                return object_storage_client.upload_part(
                    kwargs["namespace"],
                    kwargs["bucket_name"],
                    kwargs["object_name"],
                    kwargs["upload_id"],
                    kwargs["part_num"],
                    bpr,
                    **kwargs['new_kwargs'])

        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def add_parts_from_file(self, file_path):
        """
        Splits a file into parts and adds all parts to an internal list of parts to upload.  The parts will not be uploaded to the server until upload is called.

        :param string file_path: Path of file to upload in parts
        """
        with io.open(file_path, mode='rb') as file_object:
            file_object.seek(0, io.SEEK_END)
            end = file_object.tell()
            file_object.seek(0, io.SEEK_SET)
            offset = 0
            while file_object.tell() < end:
                remaining = end - offset
                self.add_part_from_file(file_path, offset=offset, size=self.part_size if remaining > self.part_size else remaining)
                offset += self.part_size
                file_object.seek(offset, io.SEEK_SET)

    def add_part_from_file(self,
                           file_path,
                           offset=0,
                           size=None,
                           part_hash=None):
        """
        Add a part to internal list of parts to upload.  The part will not be uploaded to the server until upload is called.

        :param str file_path: Path to file
        :param offset: Offset of part in file
        :param size: Size of the part in bytes
        :param part_hash: Base64 encoded MD5 hash of the part
        """
        if size is None:
            size = stat(file_path).st_size

        self.manifest["parts"].append({"file_path": file_path,
                                       "offset": offset,
                                       "size": size,
                                       "hash": part_hash})

    def abort(self, **kwargs):
        """
        Abort the multipart upload

        :param str upload_id (optional): The upload id to abort. If not provided, will attempt to abort the upload created internally by new_upload.

        :param str opc_client_request_id: (optional)
            The client request ID for tracing
        """
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()
            if 'upload_id' in kwargs:
                self.manifest["uploadId"] = kwargs['upload_id']
                kwargs.pop('upload_id')

            if not self.manifest["uploadId"]:
                raise ValueError("Cannot abort without an upload id.")

            self.object_storage_client.abort_multipart_upload(self.manifest["namespace"],
                                                              self.manifest["bucketName"],
                                                              self.manifest["objectName"],
                                                              self.manifest["uploadId"],
                                                              **kwargs)
        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def resume(self, **kwargs):
        """
       Resume uploading a multipart object to Object Storage

       :param str upload_id (optional): The upload id to resume. If not provided, will attempt to resume the upload created internally by new_upload.

       :param function progress_callback (optional):
            Callback function to receive the number of bytes uploaded since
            the last call to the callback function.

       :param str opc_client_request_id: (optional)
            The client request ID for tracing
       """
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()
            if 'upload_id' in kwargs:
                self.manifest["uploadId"] = kwargs['upload_id']
                kwargs.pop('upload_id')

            # Verify that the upload id is valid
            if self.manifest["uploadId"] is None:
                raise ValueError("Cannot resume without an upload id.")

            upload_kwargs = {}
            if 'progress_callback' in kwargs:
                upload_kwargs['progress_callback'] = kwargs['progress_callback']
                kwargs.pop('progress_callback')

            # Get parts details from object storage to see which parts didn't complete
            has_next_page = True
            while has_next_page:
                response = self.object_storage_client.list_multipart_upload_parts(self.manifest["namespace"],
                                                                                  self.manifest["bucketName"],
                                                                                  self.manifest["objectName"],
                                                                                  self.manifest["uploadId"],
                                                                                  **kwargs)
                # Update manifest with information from object storage
                parts = self.manifest["parts"]
                for part in response.data:
                    part_index = part.part_number - 1
                    if -1 < part_index < len(parts):
                        manifest_part = parts[part_index]
                        if manifest_part["size"] != part.size:
                            raise ValueError('Cannot resume upload with different part size. Parts were uploaded with a part size of {} MiB'.format(part.size / MEBIBYTE))
                        manifest_part["etag"] = part.etag
                        manifest_part["opc_md5"] = part.md5
                    elif part_index >= len(parts):
                        raise ValueError('There are more parts on the server than parts to resume, please check the upload ID.')
                has_next_page = response.has_next_page
                kwargs['page'] = response.next_page

            # Upload parts that are missing or incomplete
            self.upload(**upload_kwargs)
        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def new_upload(self, **kwargs):
        """
        Create a new multipart upload in Object Storage and add the upload
        id to the manifest

        :param str opc_client_request_id: (optional)
            The client request ID for tracing.
        """
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()

            if self.manifest['uploadId']:
                raise RuntimeError('Cannot call new_upload again once an upload has already been created.')

            request = models.CreateMultipartUploadDetails()
            request.object = self.manifest["objectName"]
            if self.content_type:
                request.content_type = self.content_type
            if self.content_language:
                request.content_language = self.content_language
            if self.content_encoding:
                request.content_encoding = self.content_encoding
            if self.metadata:
                # TODO: look into moving this into codegen for create_multipart_upload like it is for put_object
                processed_metadata = {}
                for key, value in six.iteritems(self.metadata):
                    if not key.startswith('opc-meta-'):
                        processed_metadata["opc-meta-" + key] = value
                    else:
                        processed_metadata[key] = value
                self.metadata = processed_metadata

                request.metadata = self.metadata

            client_request_id = None
            if 'opc_client_request_id' in kwargs:
                client_request_id = kwargs['opc_client_request_id']

            kwargs = {}
            if client_request_id:
                kwargs['opc_client_request_id'] = client_request_id
            if self.if_match:
                kwargs['if_match'] = self.if_match
            if self.if_none_match:
                kwargs['if_none_match'] = self.if_none_match

            # pass on SSE-C values (if any)
            kwargs.update(self.ssec_params)

            response = self.object_storage_client.create_multipart_upload(self.manifest["namespace"],
                                                                          self.manifest["bucketName"],
                                                                          request,
                                                                          **kwargs)

            self.manifest["uploadId"] = response.data.upload_id
        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def _upload_part(self, part_num, part, **kwargs):
        """
        Upload a part to Object Storage

        :param int part_num: Sequence number for where this part belongs in the
                             multipart object.
        :param dict part: Part dictionary containing the following keys: "file_path" (str), "hash" (str), "offset" (int), and "size" (int)
        :param str opc_client_request_id: (optional)
            The client request ID for tracing.
        :param function progress_callback (optional):
            Callback function to receive the number of bytes uploaded since
            the last call to the callback function.
        """
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()

            new_kwargs = {'content_md5': part["hash"]}
            if 'opc_client_request_id' in kwargs:
                new_kwargs['opc_client_request_id'] = kwargs['opc_client_request_id']

            # supply SSE-C key (if any) information to upload-part
            new_kwargs.update(self.ssec_params)

            # TODO: Calculate the hash without needing to read the file chunk twice.
            # Calculate the hash before uploading.  The hash will be used
            # to determine if the part needs to be uploaded.  It will also
            # be used to determine if there is a conflict between parts that
            # have previously been uploaded.
            if part["hash"] is None:
                part["hash"] = self.calculate_md5(part["file_path"], part["offset"], part["size"])

            retry_strategy = None
            if 'retry_strategy' in kwargs:
                retry_strategy = kwargs['retry_strategy']

            if "opc_md5" not in part:
                # Disable the retry_strategy to work around data corruption issue temporarily
                retry_strategy = None

                if retry_strategy:
                    response = retry_strategy.make_retrying_call(
                        self._upload_part_call,
                        self.object_storage_client,
                        namespace=self.manifest["namespace"],
                        bucket_name=self.manifest["bucketName"],
                        object_name=self.manifest["objectName"],
                        upload_id=self.manifest["uploadId"],
                        part_file_path=part["file_path"],
                        offset=part["offset"],
                        size=part["size"],
                        part_num=part_num,
                        new_kwargs=new_kwargs
                    )
                else:
                    # Disable the retry_strategy to work around data corruption issue temporarily
                    remaining_tries = 1
                    while remaining_tries > 0:
                        try:
                            response = self._upload_part_call(self.object_storage_client,
                                                              namespace=self.manifest["namespace"],
                                                              bucket_name=self.manifest["bucketName"],
                                                              object_name=self.manifest["objectName"],
                                                              upload_id=self.manifest["uploadId"],
                                                              part_file_path=part["file_path"],
                                                              offset=part["offset"],
                                                              size=part["size"],
                                                              part_num=part_num,
                                                              new_kwargs=new_kwargs)
                        except Exception as e:
                            if self._is_exception_retryable(e) and remaining_tries > 1:
                                remaining_tries -= 1
                            else:
                                raise
                        else:
                            break

                if response.status == 200:
                    part["etag"] = response.headers['etag']
                    part["opc_md5"] = str(response.headers['opc-content-md5'])

            if 'progress_callback' in kwargs:
                kwargs['progress_callback'](part["size"])

        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def _upload_stream_part(self, part_num, part_bytes, **kwargs):
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()

            m = hashlib.md5()
            m.update(part_bytes)

            new_kwargs = {'content_md5': base64.b64encode(m.digest()).decode("utf-8")}
            if 'opc_client_request_id' in kwargs:
                new_kwargs['opc_client_request_id'] = kwargs['opc_client_request_id']

            remaining_tries = self.max_retries
            while remaining_tries > 0:
                try:
                    response = self.object_storage_client.upload_part(
                        self.manifest["namespace"],
                        self.manifest["bucketName"],
                        self.manifest["objectName"],
                        self.manifest["uploadId"],
                        part_num + 1,  # Internally this is 0-based but object storage is 1-based
                        part_bytes,
                        **new_kwargs
                    )
                except Exception as e:
                    if self._is_exception_retryable(e) and remaining_tries > 1:
                        remaining_tries -= 1
                    else:
                        if 'shared_dict' in kwargs:
                            kwargs['shared_dict']['should_continue'] = False
                            kwargs['shared_dict']['exceptions'].put(e)
                        raise
                else:
                    break

            if response.status == 200:
                self.manifest['parts'].append({
                    'etag': response.headers['etag'],
                    'opc_md5': str(response.headers['opc-content-md5']),
                    'part_num': part_num
                })

            if 'progress_callback' in kwargs:
                kwargs['progress_callback'](len(part_bytes))
        finally:
            if 'semaphore' in kwargs:
                kwargs['semaphore'].release()

            # Reset the object storage client timeout regardless of the result of above operations
            self._reset_object_storage_client_timeout()

    def upload(self, **kwargs):
        """
        Upload an object to Object Storage

        :param function progress_callback (optional):
            Callback function to receive the number of bytes uploaded since
            the last call to the callback function.
        :param str opc_client_request_id: (optional)
            The client request ID for tracing.
        """
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()

            if self.manifest["uploadId"] is None:
                raise RuntimeError('Cannot call upload before initializing an upload using new_upload.')

            pool = Pool(processes=self.parallel_process_count)
            pool.map(lambda part_tuple: self._upload_part(part_num=part_tuple[0] + 1, part=part_tuple[1], **kwargs),
                     enumerate(self.manifest["parts"]))
        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def upload_stream(self, stream_ref, **kwargs):
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()

            if self.manifest["uploadId"] is None:
                raise RuntimeError('Cannot call upload before initializing an upload using new_upload.')

            # The pool of work we have available, and the sempahore to gate work into the pool (since just submitting
            # work to the pool doesn't block on the number of processes available to do work in the pool)
            pool = Pool(processes=self.parallel_process_count)
            semaphore = Semaphore(self.parallel_process_count)

            # A dict which will be shared between the threads in our pool (this would not work as-is with processes) but
            # we use threads via multiprocessing.dummy. If we use processes, then a Manager would likely be needed for this.
            #
            # should_continue will only ever be set to False by _upload_stream_part so not too worried if we have multiple
            # writers
            #
            # Queue should be thread safe (though for tracking the exceptions, order doesn't strictly matter)
            shared_dict = {'should_continue': True, 'exceptions': Queue()}

            part_counter = 0

            apply_async_kwargs = kwargs.copy()
            apply_async_kwargs['semaphore'] = semaphore
            apply_async_kwargs['shared_dict'] = shared_dict

            # We pull data from the stream until there is no more
            keep_reading = True
            while keep_reading:
                if six.PY3 and hasattr(stream_ref, 'buffer'):
                    read_bytes = stream_ref.buffer.read(self.part_size)
                else:
                    read_bytes = stream_ref.read(self.part_size)

                semaphore.acquire()

                if len(read_bytes) != 0:
                    pool.apply_async(self._upload_stream_part, (part_counter, read_bytes), apply_async_kwargs)
                    part_counter += 1

                keep_reading = (len(read_bytes) == self.part_size) and shared_dict['should_continue']

            # If we're here we've either sent off all the work we needed to (and so are waiting on remaining bits to finish)
            # or we terminated early because of an exception in one of our uploads. In either case, close off the pool to
            # any more work and let the remaining work finish gracefully
            pool.close()
            pool.join()

            # If we had at least one exception then throw out an error to indicate failure
            if not shared_dict['exceptions'].empty():
                raise MultipartUploadError(error_causes_queue=shared_dict['exceptions'])

            # Because we processed in parallel, the parts in the manifest may be out of order. Re-order them based on the part number
            # because commit assumes that they are ordered
            self.manifest['parts'].sort(key=lambda part: part['part_num'])
        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    def commit(self, **kwargs):
        """
        Commit the multipart upload.

        :return: A Response object with data of type None
        :rtype: None
        """
        try:
            # Set the upload manager timeout on object storage client
            self._set_multipart_upload_timeout()

            if self.manifest["uploadId"] is None:
                raise RuntimeError('Cannot call commit before initializing an upload using new_upload or resuming an upload using resume.')

            commit_details = models.CommitMultipartUploadDetails()

            # Determine which parts to commit and which parts to exclude.
            parts_to_commit = []
            parts_to_exclude = []
            for partNum, part in enumerate(self.manifest["parts"]):
                detail = models.CommitMultipartUploadPartDetails()
                detail.part_num = partNum + 1
                if "etag" in part:
                    detail.etag = part["etag"]
                    parts_to_commit.append(detail)
                else:
                    parts_to_exclude.append(partNum + 1)

            commit_details.parts_to_commit = parts_to_commit
            commit_details.parts_to_exclude = parts_to_exclude

            # Commit the multipart upload
            response = self.object_storage_client.commit_multipart_upload(self.manifest["namespace"],
                                                                          self.manifest["bucketName"],
                                                                          self.manifest["objectName"],
                                                                          self.manifest["uploadId"],
                                                                          commit_details,
                                                                          **kwargs)

            return response
        # Reset the object storage client timeout regardless of the result of above operations
        finally:
            self._reset_object_storage_client_timeout()

    # The following two methods are called at the beginning and at the end of every method for upload manager
    # This is not the best way to go around the timeout issue, but it is a workaround for now
    def _set_multipart_upload_timeout(self):
        self.object_storage_client.base_client.timeout = MULTIPART_UPLOAD_TIMEOUT

    def _reset_object_storage_client_timeout(self):
        self.object_storage_client.base_client.timeout = self._object_storage_client_timeout
