#!/usr/bin/python
#
# Copyright 2018-2020 Polyaxon, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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 os

from concurrent import futures
from io import BytesIO

from botocore.exceptions import ClientError

from polyaxon.connections.aws import base
from polyaxon.connections.aws.base import AWSService
from polyaxon.exceptions import (
    PolyaxonPathException,
    PolyaxonSchemaError,
    PolyaxonStoresException,
)
from polyaxon.logger import logger
from polyaxon.parser import parser
from polyaxon.stores.base_store import StoreMixin
from polyaxon.utils.date_utils import file_modified_since
from polyaxon.utils.path_utils import (
    append_basename,
    check_dirname_exists,
    get_files_in_path_context,
)
from polyaxon.utils.string_utils import force_bytes


class S3Service(AWSService, StoreMixin):
    """
    S3 store Service using Boto3.
    """

    RESOURCE_TYPE = "s3"

    @staticmethod
    def parse_s3_url(s3_url):
        """
        Parses and validates an S3 url.

        Returns:
             tuple(bucket_name, key).
        """
        try:
            spec = parser.parse_s3_path(s3_url)
            return spec.bucket, spec.key
        except PolyaxonSchemaError as e:
            raise PolyaxonStoresException("Connection error: %s" % e) from e

    @staticmethod
    def check_prefix_format(prefix, delimiter):
        if not delimiter or not prefix:
            return prefix
        return prefix + delimiter if prefix[-1] != delimiter else prefix

    def check_bucket(self, bucket_name):
        """
        Checks if a buckete exists.

        Args:
            bucket_name: `str`. Name of the bucket
        """
        try:
            self.connection.head_bucket(Bucket=bucket_name)
            return True
        except ClientError as e:
            logger.info(e.response["Error"]["Message"])
            return False

    def get_bucket(self, bucket_name):
        """
        Gets a bucket by name.

        Args:
            bucket_name: `str`. Name of the bucket
        """
        return self.resource.Bucket(bucket_name)

    def ls(self, path):
        (bucket_name, key) = self.parse_s3_url(path)
        results = self.list(bucket_name=bucket_name, prefix=key)
        return {"files": results["keys"], "dirs": results["prefixes"]}

    def list(
        self,
        bucket_name,
        prefix="",
        delimiter="/",
        page_size=None,
        max_items=None,
        keys=True,
        prefixes=True,
    ):
        """
        Lists prefixes and contents in a bucket under prefix.

        Args:
            bucket_name: `str`. the name of the bucket
            prefix: `str`. a key prefix
            delimiter: `str`. the delimiter marks key hierarchy.
            page_size: `str`. pagination size
            max_items: `int`. maximum items to return
            keys: `bool`. if it should include keys
            prefixes: `boll`. if it should include prefixes
        """
        config = {"PageSize": page_size, "MaxItems": max_items}

        legacy_api = base.get_legacy_api(legacy_api=self._aws_legacy_api)

        if legacy_api:
            paginator = self.connection.get_paginator("list_objects")
        else:
            paginator = self.connection.get_paginator("list_objects_v2")

        prefix = self.check_prefix_format(prefix=prefix, delimiter=delimiter)
        response = paginator.paginate(
            Bucket=bucket_name,
            Prefix=prefix,
            Delimiter=delimiter,
            PaginationConfig=config,
        )

        def get_keys(contents):
            list_keys = []
            for cont in contents:
                list_keys.append((cont["Key"][len(prefix) :], cont.get("Size")))

            return list_keys

        def get_prefixes(page_prefixes):
            list_prefixes = []
            for pref in page_prefixes:
                list_prefixes.append(pref["Prefix"][len(prefix) : -1])
            return list_prefixes

        results = {"keys": [], "prefixes": []}
        for page in response:
            if prefixes:
                results["prefixes"] += get_prefixes(page.get("CommonPrefixes", []))
            if keys:
                results["keys"] += get_keys(page.get("Contents", []))

        return results

    def list_prefixes(
        self, bucket_name, prefix="", delimiter="", page_size=None, max_items=None
    ):
        """
        Lists prefixes in a bucket under prefix

        Args:
            bucket_name: `str`. the name of the bucket
            prefix: `str`. a key prefix
            delimiter: `str`. the delimiter marks key hierarchy.
            page_size: `int`. pagination size
            max_items: `int`. maximum items to return
        """
        results = self.list(
            bucket_name=bucket_name,
            prefix=prefix,
            delimiter=delimiter,
            page_size=page_size,
            max_items=max_items,
            keys=False,
            prefixes=True,
        )
        return results["prefixes"]

    def list_keys(
        self, bucket_name, prefix="", delimiter="", page_size=None, max_items=None
    ):
        """
        Lists keys in a bucket under prefix and not containing delimiter

        Args:
            bucket_name: `str`. the name of the bucket
            prefix: `str`. a key prefix
            delimiter: `str`. the delimiter marks key hierarchy.
            page_size: `int`. pagination size
            max_items: `int`. maximum items to return
        """
        results = self.list(
            bucket_name=bucket_name,
            prefix=prefix,
            delimiter=delimiter,
            page_size=page_size,
            max_items=max_items,
            keys=True,
            prefixes=False,
        )
        return results["keys"]

    def check_key(self, key, bucket_name=None):
        """
        Checks if a key exists in a bucket

        Args:
            key: `str`. S3 key that will point to the file
            bucket_name: `str`. Name of the bucket in which the file is stored
        """
        if not bucket_name:
            (bucket_name, key) = self.parse_s3_url(key)

        try:
            self.connection.head_object(Bucket=bucket_name, Key=key)
            return True
        except ClientError as e:
            logger.info(e.response["Error"]["Message"])
            return False

    def get_key(self, key, bucket_name=None):
        """
        Returns a boto3.s3.Object

        Args:
            key: `str`. the path to the key.
            bucket_name: `str`. the name of the bucket.
        """
        if not bucket_name:
            (bucket_name, key) = self.parse_s3_url(key)

        try:
            obj = self.resource.Object(bucket_name, key)
            obj.load()
            return obj
        except Exception as e:
            raise PolyaxonStoresException("Connection error: %s" % e) from e

    def read_key(self, key, bucket_name=None):
        """
        Reads a key from S3

        Args:
            key: `str`. S3 key that will point to the file.
            bucket_name: `str`. Name of the bucket in which the file is stored.
        """

        obj = self.get_key(key, bucket_name)
        return obj.get()["Body"].read().decode("utf-8")

    def upload_bytes(
        self,
        bytes_data,
        key,
        bucket_name=None,
        overwrite=False,
        encrypt=False,
        acl=None,
    ):
        """
        Uploads bytes to S3

        This is provided as a convenience to drop a string in S3. It uses the
        boto infrastructure to ship a file to s3.

        Args:
            bytes_data: `bytes`. bytes to set as content for the key.
            key: `str`. S3 key that will point to the file.
            bucket_name: `str`. Name of the bucket in which to store the file.
            overwrite: `bool`. A flag to decide whether or not to overwrite the key
                if it already exists.
            encrypt: `bool`. If True, the file will be encrypted on the server-side
                by S3 and will be stored in an encrypted form while at rest in S3.
            acl: `str`. ACL to use for uploading, e.g. "public-read".
        """
        if not bucket_name:
            (bucket_name, key) = self.parse_s3_url(key)

        if not overwrite and self.check_key(key, bucket_name):
            raise ValueError("The key {key} already exists.".format(key=key))

        extra_args = {}
        if encrypt:
            extra_args["ServerSideEncryption"] = self.ENCRYPTION
        if acl:
            extra_args["ACL"] = acl

        filelike_buffer = BytesIO(bytes_data)

        self.connection.upload_fileobj(
            filelike_buffer, bucket_name, key, ExtraArgs=extra_args
        )

    def upload_string(
        self,
        string_data,
        key,
        bucket_name=None,
        overwrite=False,
        encrypt=False,
        acl=None,
        encoding="utf-8",
    ):
        """
        Uploads a string to S3.

        This is provided as a convenience to drop a string in S3. It uses the
        boto infrastructure to ship a file to s3.

        Args:
            string_data: `str`. string to set as content for the key.
            key: `str`. S3 key that will point to the file.
            bucket_name: `str`. Name of the bucket in which to store the file.
            overwrite: `bool`. A flag to decide whether or not to overwrite the key
                if it already exists.
            encrypt: `bool`. If True, the file will be encrypted on the server-side
                by S3 and will be stored in an encrypted form while at rest in S3.
            acl: `str`. ACL to use for uploading, e.g. "public-read".
            encoding: `str`. Encoding to use.
        """
        self.upload_bytes(
            force_bytes(string_data, encoding=encoding),
            key=key,
            bucket_name=bucket_name,
            overwrite=overwrite,
            encrypt=encrypt,
            acl=acl,
        )

    def upload_file(
        self,
        filename,
        key,
        bucket_name=None,
        overwrite=False,
        encrypt=False,
        acl=None,
        use_basename=True,
    ):
        """
        Uploads a local file to S3.

        Args:
            filename: `str`. name of the file to upload.
            key: `str`. S3 key that will point to the file.
            bucket_name: `str`. Name of the bucket in which to store the file.
            overwrite: `bool`. A flag to decide whether or not to overwrite the key
                if it already exists. If replace is False and the key exists, an
                error will be raised.
            encrypt: `bool`. If True, the file will be encrypted on the server-side
                by S3 and will be stored in an encrypted form while at rest in S3.
            acl: `str`. ACL to use for uploading, e.g. "public-read".
            use_basename: `bool`. whether or not to use the basename of the filename.
        """
        if not bucket_name:
            bucket_name, key = self.parse_s3_url(key)

        if use_basename:
            key = append_basename(key, filename)

        if not overwrite and self.check_key(key, bucket_name):
            raise PolyaxonStoresException("The key {} already exists.".format(key))

        extra_args = {}
        if encrypt:
            extra_args["ServerSideEncryption"] = self.ENCRYPTION
        if acl:
            extra_args["ACL"] = acl

        self.connection.upload_file(filename, bucket_name, key, ExtraArgs=extra_args)

    def download_file(self, key, local_path, bucket_name=None, use_basename=True):
        """
        Download a file from S3.

        Args:
            key: `str`. S3 key that will point to the file.
            local_path: `str`. the path to download to.
            bucket_name: `str`. Name of the bucket in which to store the file.
            use_basename: `bool`. whether or not to use the basename of the key.
        """
        if not bucket_name:
            bucket_name, key = self.parse_s3_url(key)

        local_path = os.path.abspath(local_path)

        if use_basename:
            local_path = append_basename(local_path, key)

        try:
            check_dirname_exists(local_path)
        except PolyaxonPathException as e:
            raise PolyaxonStoresException("Connection error: %s" % e) from e

        try:
            self.connection.download_file(bucket_name, key, local_path)
        except ClientError as e:
            raise PolyaxonStoresException("Connection error: %s" % e) from e

    def upload_dir(
        self,
        dirname,
        key,
        bucket_name=None,
        overwrite=False,
        encrypt=False,
        acl=None,
        use_basename=True,
        workers=0,
        last_time=None,
    ):
        """
        Uploads a local directory to S3.

        Args:
            dirname: `str`. name of the directory to upload.
            key: `str`. S3 key that will point to the file.
            bucket_name: `str`. Name of the bucket in which to store the file.
            overwrite: `bool`. A flag to decide whether or not to overwrite the key
                if it already exists. If replace is False and the key exists, an
                error will be raised.
            encrypt: `bool`. If True, the file will be encrypted on the server-side
                by S3 and will be stored in an encrypted form while at rest in S3.
            acl: `str`. ACL to use for uploading, e.g. "public-read".
            use_basename: `bool`. whether or not to use the basename of the directory.
            last_time: `datetime`. If provided will only upload the file if changed after last_time.
        """
        if not bucket_name:
            bucket_name, key = self.parse_s3_url(key)

        if use_basename:
            key = append_basename(key, dirname)

        pool, future_results = self.init_pool(workers)

        # Turn the path to absolute paths
        dirname = os.path.abspath(dirname)
        with get_files_in_path_context(dirname) as files:
            for f in files:

                # If last time is provided we check if we should re-upload the file
                if last_time and not file_modified_since(
                    filepath=f, last_time=last_time
                ):
                    continue

                file_key = os.path.join(key, os.path.relpath(f, dirname))
                future_results = self.submit_pool(
                    workers=workers,
                    pool=pool,
                    future_results=future_results,
                    fn=self.upload_file,
                    filename=f,
                    key=file_key,
                    bucket_name=bucket_name,
                    overwrite=overwrite,
                    encrypt=encrypt,
                    acl=acl,
                    use_basename=False,
                )
        if workers:
            futures.wait(future_results)
            self.close_pool(pool=pool)

    def download_dir(
        self,
        key: str,
        local_path: str,
        bucket_name: str = None,
        use_basename: bool = True,
        workers: int = 0,
    ):
        """
        Download a directory from S3.

        Args:
            key: `str`. S3 key that will point to the file.
            local_path: `str`. the path to download to.
            bucket_name: `str`. Name of the bucket in which to store the file.
            use_basename: `bool`. whether or not to use the basename of the key.
            workers: number of workers threads to use for parallel execution.
        """
        if not bucket_name:
            bucket_name, key = self.parse_s3_url(key)

        local_path = os.path.abspath(local_path)

        if use_basename:
            local_path = append_basename(local_path, key)

        try:
            check_dirname_exists(local_path, is_dir=True)
        except PolyaxonPathException:
            os.makedirs(local_path)

        results = self.list(bucket_name=bucket_name, prefix=key, delimiter="/")

        # Create directories
        for prefix in sorted(results["prefixes"]):
            direname = os.path.join(local_path, prefix)
            prefix = os.path.join(key, prefix)
            # Download files under
            self.download_dir(
                key=prefix,
                local_path=direname,
                bucket_name=bucket_name,
                use_basename=False,
            )

        pool, future_results = self.init_pool(workers)

        # Download files
        for file_key in results["keys"]:
            file_key = file_key[0]
            filename = os.path.join(local_path, file_key)
            file_key = os.path.join(key, file_key)
            future_results = self.submit_pool(
                workers=workers,
                pool=pool,
                future_results=future_results,
                fn=self.download_file,
                key=file_key,
                local_path=filename,
                bucket_name=bucket_name,
                use_basename=False,
            )

        if workers:
            futures.wait(future_results)
            self.close_pool(pool=pool)

    def delete(self, key, bucket_name=None, workers=0):
        if not bucket_name:
            (bucket_name, key) = self.parse_s3_url(key)

        results = self.list(bucket_name=bucket_name, prefix=key, delimiter="/")

        if not any([results["prefixes"], results["keys"]]):
            self.delete_file(key=key, bucket_name=bucket_name)

        # Delete directories
        for prefix in sorted(results["prefixes"]):
            prefix = os.path.join(key, prefix)
            # Download files under
            self.delete(key=prefix, bucket_name=bucket_name)

        pool, future_results = self.init_pool(workers)

        # Delete files
        for file_key in results["keys"]:
            file_key = file_key[0]
            file_key = os.path.join(key, file_key)
            future_results = self.submit_pool(
                workers=workers,
                pool=pool,
                future_results=future_results,
                fn=self.delete_file,
                key=file_key,
                bucket_name=bucket_name,
            )
        if workers:
            futures.wait(future_results)
            self.close_pool(pool=pool)

    def delete_file(self, key, bucket_name=None):
        if not bucket_name:
            (bucket_name, key) = self.parse_s3_url(key)
        try:
            obj = self.resource.Object(bucket_name, key)
            obj.delete()
        except ClientError as e:
            raise PolyaxonStoresException("Connection error: %s" % e) from e
