import os

from boto.s3.connection import S3Connection
import boto3
from boto3.s3.transfer import S3Transfer
from garcon_contrib.aws.utils import garcon_s3


def download_boto3(s3_path, local_path):
    """Download S3 object(s) to local directory

    If number of objects is more than 1 in the s3_path, the names of the
    downloaded files will be the local_path plus a suffix of the file name on
    S3. Otherwise, downloaded file will be renamed to local_path completely.

    Args:
        s3_path (str): prefix of s3 object(s). Eg. s3://bucket/file - will get
         s3://bucket/file_part1.gz, s3://bucket/file_part2.gz.....
        local_path (str): local directory.

    Returns:
        list: downloaded files

    Raises:
        Exception: if s3_path is not found.
    """

    client = boto3.client('s3')
    transfer = S3Transfer(client)
    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path)
    responses = client.list_objects(Bucket=bucket, Prefix=bucket_path)

    if 'Contents' not in responses:
        raise Exception('Object does not exist: {}'.format(s3_path))
    response_keys = [item.get('Key') for item in responses['Contents']
                     if os.path.basename(item.get('Key'))]

    # only one file to download.
    if len(response_keys) == 1:
        key = response_keys.pop(0)
        transfer.download_file(bucket, key, local_path)
        return [key]

    # multiple files to download
    downloaded_files = []
    for key in response_keys:
        local_filename = '{}_{}'.format(local_path, key)
        transfer.download_file(bucket, key, local_filename)
        downloaded_files.append(local_filename)
    return downloaded_files


def upload_boto3(local_path, s3_path):
    """Upload S3 object(s) to designated s3 path

    Args:
        local_path (str): full local path to the file
        s3_path (str): full destination s3 path
    """

    client = boto3.client('s3')
    transfer = S3Transfer(client)
    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path)

    transfer.upload_file(local_path, bucket, bucket_path)


def delete_boto3(s3_path):
    """Delete an S3 object.

    Args:
        s3_path (str): full path to the S3 object
    """
    client = boto3.client('s3')
    bucket, key_path = garcon_s3.extract_bucket_path(s3_path)
    client.delete_object(Bucket=bucket, Key=key_path)


def download(s3_path, local_path):
    """Download S3 object(s) to local directory

    Args:
        s3_path (str): prefix of s3 object(s). Eg. s3://bucket/file - will get
         s3://bucket/file_part1.gz, s3://bucket/file_part2.gz.....
        local_path (str): local directory
    """

    s3_connection = S3Connection()
    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path)
    bucket_object = s3_connection.get_bucket(bucket)

    for key in bucket_object.list(bucket_path):
        key.get_contents_to_filename(local_path)


def upload_raw_file_to_s3(source_path, s3_path):
    """Upload local file to s3

    Args:
        source_path (str): full path to local path
        s3_path (str): full S3 destination path
    """

    s3_connection = S3Connection()
    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_path)
    bucket_object = s3_connection.get_bucket(bucket)

    key_object = bucket_object.new_key(bucket_path)
    key_object.set_contents_from_filename(source_path, None, True)