# Copyright 2019-2020 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.
"""This module contains Enums and helper methods related to S3."""
from __future__ import print_function, absolute_import

import logging

from six.moves.urllib.parse import urlparse
from sagemaker.session import Session

logger = logging.getLogger("sagemaker")

SESSION_V2_RENAME_MESSAGE = (
    "Parameter 'session' will be renamed to 'sagemaker_session' in SageMaker Python SDK v2."
)


def _session_v2_rename_warning(session):
    """
    Args:
        session (sagemaker.session.Session):
    """
    if session is not None:
        logger.warning(SESSION_V2_RENAME_MESSAGE)


def parse_s3_url(url):
    """Returns an (s3 bucket, key name/prefix) tuple from a url with an s3
    scheme.
    Args:
        url (str):
    Returns:
        tuple: A tuple containing:
            str: S3 bucket name str: S3 key
    """
    parsed_url = urlparse(url)
    if parsed_url.scheme != "s3":
        raise ValueError("Expecting 's3' scheme, got: {} in {}.".format(parsed_url.scheme, url))
    return parsed_url.netloc, parsed_url.path.lstrip("/")


class S3Uploader(object):
    """Contains static methods for uploading directories or files to S3."""

    @staticmethod
    def upload(local_path, desired_s3_uri, kms_key=None, session=None):
        """Static method that uploads a given file or directory to S3.

        Args:
            local_path (str): Path (absolute or relative) of local file or directory to upload.
            desired_s3_uri (str): The desired S3 location to upload to. It is the prefix to
                which the local filename will be added.
            kms_key (str): The KMS key to use to encrypt the files.
            session (sagemaker.session.Session): Session object which
                manages interactions with Amazon SageMaker APIs and any other
                AWS services needed. If not specified, the estimator creates one
                using the default AWS configuration chain.

        Returns:
            The S3 uri of the uploaded file(s).

        """
        if session is not None:
            _session_v2_rename_warning(session)

        sagemaker_session = session or Session()
        bucket, key_prefix = parse_s3_url(url=desired_s3_uri)
        if kms_key is not None:
            extra_args = {"SSEKMSKeyId": kms_key}
        else:
            extra_args = None

        return sagemaker_session.upload_data(
            path=local_path, bucket=bucket, key_prefix=key_prefix, extra_args=extra_args
        )

    @staticmethod
    def upload_string_as_file_body(body, desired_s3_uri=None, kms_key=None, session=None):
        """Static method that uploads a given file or directory to S3.

        Args:
            body (str): String representing the body of the file.
            desired_s3_uri (str): The desired S3 uri to upload to.
            kms_key (str): The KMS key to use to encrypt the files.
            session (sagemaker.session.Session): AWS session to use. Automatically
                generates one if not provided.

        Returns:
            str: The S3 uri of the uploaded file(s).

        """
        if session is not None:
            _session_v2_rename_warning(session)

        sagemaker_session = session or Session()
        bucket, key = parse_s3_url(desired_s3_uri)

        sagemaker_session.upload_string_as_file_body(
            body=body, bucket=bucket, key=key, kms_key=kms_key
        )

        return desired_s3_uri


class S3Downloader(object):
    """Contains static methods for downloading directories or files from S3."""

    @staticmethod
    def download(s3_uri, local_path, kms_key=None, session=None):
        """Static method that downloads a given S3 uri to the local machine.

        Args:
            s3_uri (str): An S3 uri to download from.
            local_path (str): A local path to download the file(s) to.
            kms_key (str): The KMS key to use to decrypt the files.
            session (sagemaker.session.Session): Session object which
                manages interactions with Amazon SageMaker APIs and any other
                AWS services needed. If not specified, the estimator creates one
                using the default AWS configuration chain.

        """
        if session is not None:
            _session_v2_rename_warning(session)

        sagemaker_session = session or Session()
        bucket, key_prefix = parse_s3_url(url=s3_uri)
        if kms_key is not None:
            extra_args = {"SSECustomerKey": kms_key}
        else:
            extra_args = None

        sagemaker_session.download_data(
            path=local_path, bucket=bucket, key_prefix=key_prefix, extra_args=extra_args
        )

    @staticmethod
    def read_file(s3_uri, session=None):
        """Static method that returns the contents of an s3 uri file body as a string.

        Args:
            s3_uri (str): An S3 uri that refers to a single file.
            session (sagemaker.session.Session): AWS session to use. Automatically
                generates one if not provided.

        Returns:
            str: The body of the file.

        """
        if session is not None:
            _session_v2_rename_warning(session)

        sagemaker_session = session or Session()
        bucket, key_prefix = parse_s3_url(url=s3_uri)

        return sagemaker_session.read_s3_file(bucket=bucket, key_prefix=key_prefix)

    @staticmethod
    def list(s3_uri, session=None):
        """Static method that lists the contents of an S3 uri.

        Args:
            s3_uri (str): The S3 base uri to list objects in.
            session (sagemaker.session.Session): AWS session to use. Automatically
                generates one if not provided.

        Returns:
            [str]: The list of S3 URIs in the given S3 base uri.

        """
        if session is not None:
            _session_v2_rename_warning(session)

        sagemaker_session = session or Session()
        bucket, key_prefix = parse_s3_url(url=s3_uri)

        file_keys = sagemaker_session.list_s3_files(bucket=bucket, key_prefix=key_prefix)
        return ["s3://{}/{}".format(bucket, file_key) for file_key in file_keys]
