"""
AWS Simple Storage Service (S3)
===============================

All the tasks that are related to the AWS S3.
"""

import os
import subprocess

import boto3
from botocore.exceptions import ClientError
from garcon import task

from garcon_contrib.aws.utils import garcon_s3 as s3_utils

import smart_open


@task.decorate()
def remove_files_from_path(activity, path, return_deleted_files):
    """Removes all the objects from a given bucket path.
    Deleted objects include anything that starts with the path.

    Args:
        activity (ActivityWorker): the swf activity worker.
        path (str): S3 URL formatted like 's3://bucket_name/folder1/folder2/'
        return_deleted_files (bool): whether return deleted file names or not

    Return:
        dict: One element dict with a list of deleted files.
    """

    bucket_name, bucket_path = s3_utils.extract_bucket_path(path)
    s3 = boto3.client('s3')
    paginator = s3.get_paginator('list_objects_v2')

    operation_parameters = {'Bucket': bucket_name, 'Prefix': bucket_path}
    page_iterator = paginator.paginate(**operation_parameters)

    deleted = []
    errors = {}
    for page in page_iterator:
        if 'Contents' in page:
            for obj in page['Contents']:
                key = obj['Key']
                try:
                    s3.delete_object(Bucket=bucket_name, Key=obj['Key'])
                    deleted.append(key)
                except ClientError as e:
                    errors[key] = str(e)

    if not return_deleted_files:
        return_deleted_files = False

    if errors:
        activity.logger.error(
            'Error while deleting objects from {path}'.format(path=path))
        raise Exception(f'Not all files in s3 have been removed: {errors}')

    if deleted:
        activity.logger.info(
            'Deleted all S3 objects from {path}'.format(path=path))

    if return_deleted_files is True:
        return {
            's3.files_removed': deleted
        }

    return {}


@task.decorate()
def create_object(activity, path, content):
    """Create file in path with content

    Create file in path with content. If file already exist, content will
    be replaced.

    Args:
        activity (ActivityWorker): the swf activity worker
        path (str): S3 URL formatted like 's3://bucket_name/folder1/file'
        content (str): content of s3 file
    """

    assert not path[-1:] == '/', 'S3 path should point to a file'
    bucket_name, bucket_path = s3_utils.extract_bucket_path(path)
    s3 = boto3.client('s3')
    s3.put_object(Bucket=bucket_name, Key=bucket_path, Body=content)


@task.decorate(timeout=7200)
def join_s3_objects(activity, source_s3_path, destination_s3_path):
    """Join S3 object

    Merge S3 objects in a source bucket/key and write it directly
    to a destination S3 bucket/key. Following are the examples:

    source_s3_path:
    's3://test_bucket/emr/soundscan/2015-04-27/CANADA_PRODUCT/'

    destination_s3_path is:
    's3://test_bucket/emr/soundscan/2015-04-27/merged/CANADA_PRODUCT/CANADA_PRODUCT.txt.gz'

    Args:
        activity (ActivityWorker): the swf activity worker
        source_s3_path (str): source path to source s3 key.
        destination_s3_path (str): destination path to merged file

    Returns:
        command (str): command ran
    """

    source_bucket, source_prefix = s3_utils.extract_bucket_path(source_s3_path)
    s3_resource = boto3.resource('s3')
    bucket = s3_resource.Bucket(source_bucket)

    file_objects = (
        ' <(aws s3 cp {source_s3_path} -) '.format(
            source_s3_path='s3://{bucket_name}/{key}'.format(
                bucket_name=source_bucket,
                key=key.key))
        for key in bucket.objects.filter(Prefix=source_prefix))

    upload_command = 'aws s3 cp - {destination_s3_path}'.format(
        destination_s3_path=destination_s3_path)

    command = 'cat {copy} | {upload}'.format(
        copy=''.join(file_objects),
        upload=upload_command)

    subprocess.call(['bash', '-c', command])


@task.decorate(timeout=7200)
def stream_join_s3_objects(activity, source_s3_path, destination_s3_path):
    """Join S3 objects into one object

    Args:
        activity (ActivityWorker): the activity worker.
        source_s3_path (str): folder where source data is located.
        destination_s3_path (str): folder where final report file is located.
    """

    source_bucket, source_key = s3_utils.extract_bucket_path(source_s3_path)
    s3_resource = boto3.resource('s3')
    bucket = s3_resource.Bucket(source_bucket)
    source_key = source_key.rstrip('/')

    with smart_open.smart_open(destination_s3_path, 'wb') as fout:
        for obj in bucket.objects.filter(Prefix=source_key):
            if os.path.basename(obj.key):
                for line in smart_open.smart_open(
                        's3://{}/{}'.format(source_bucket, obj.key)):
                    fout.write(line)