import os
import logging
from io import BytesIO
from datetime import date, timedelta, datetime
import boto3
from boto3.s3.transfer import TransferConfig, S3Transfer
from botocore.exceptions import ClientError
from zenyatta.aws import boto_client, boto_resource
import zenyatta


def get_base_s3_path(key: str, ts_nodash: str, conn_id: str) -> str:
    """
    wrap logic around the path for a task to use in s3
    :param key: zenyatta path for the final s3 key
    :param ts_nodash: keyward argument passed to TaskInstances which represents a timestamp
    :param connection: connection object to inspect metadata for database
    :return:
    """
    """
    there exists connections like this: justintv_prod-20170531T080000
    and those connections will totally munge the s3 output path to
    include the timestamp twice like this:
    dbsnapshots/20170531T080000/justintv_prod-20170531T080000/oauth_categories.csv
    """
    bad_string = "-{}".format(ts_nodash)
    if bad_string in conn_id:
        conn_id = conn_id.replace(bad_string, '')
    return "{key}/{ts_nodash}/{conn_id}".format(**locals())


def ensure_s3_bucket_exists(s3_meta: dict, conn_id: str, **kwargs) -> bool:
    s3_client, _ = boto_client('s3', s3_meta['role_arn'])
    try:
        s3_key = "{}/pipeline.start".format(get_base_s3_path(s3_meta['s3_key'], kwargs.get('ts_nodash'), conn_id))
        resp = s3_client.put_object(Bucket=s3_meta['bucket'], Key=s3_key, Body=BytesIO(b'pipeline starting'))
        if resp['ResponseMetadata']['HTTPStatusCode'] == 200:
            return True
        else:
            return False
    except ClientError as e:
        return False


def set_expiration_on_s3_object(bucket: boto3.session.Session, key: str, ttl: int) -> None:
    s3_resource, _ = boto_resource('s3', role_arn)
    bucket = s3_resource.Bucket(bucket)
    obj = bucket.Object(key=key)
    expiration_date = date.today() + timedelta(days=ttl)
    obj.put(Expires=str(expiration_date))


def check_s3_object_exist(bucket: str, role_arn: str, prefix: str) -> bool:
    """
    :param bucket: s3 bucket
    :param role_arn
    :param prefix: full object path after bucket name
    """
    s3_resource, _ = boto_resource('s3', role_arn)
    try:
        s3_bucket = s3_resource.Bucket(bucket)
        obj_list = list(s3_bucket.objects.filter(Prefix=prefix))
        return True if len(obj_list) > 0 else False
    except ClientError as e:
        return False


def upload_file_to_s3(bucket: str, file_path: str, object_name: str, role_arn: str, is_dir=False) -> bool:
    """
    upload a single file or a folder from local to s3
    :param bucket: s3 bucket
    :param file_path: full path of the file or the dir name
    :param object_name: full name of s3 object or s3 dir object
    :param role_arn
    :param is_dir: default False for single file upload
    :return:
    """
    """
    setup config to allow for 128 megabyte sized chunks to upload.
    it allows files up to 1.2T for this transfer.
    """
    s3_client, _ = boto_client('s3', role_arn, 'us-west-2', 's3v4')
    config = TransferConfig(multipart_chunksize=(128 * 1024 * 1024), max_concurrency=3)
    transfer = S3Transfer(client=s3_client, config=config)
    try:
        if not is_dir:
            results = transfer.upload_file(file_path, bucket, object_name)
            logging.info("uploaded object from {file_path} {object_name} to "
                         "{bucket} {object_name}".format(**locals()))
        else:
            files = os.listdir(file_path)
            for f in files:
                results = transfer.upload_file(file_path + "/" + f, bucket, object_name + "/" + f)
                logging.info("uploaded object from {file_path} {object_name} to "
                             "{bucket} {object_name}".format(**locals()))
        return True
    except ClientError as e:
        logging.info("failed to upload object from {file_path} {object_name} to "
                     "{bucket} {object_name}.".format(**locals()) +
                     "{}".format(str(e)))
        return False


def download_objects_from_s3(bucket: str, object_name: str, file_path: str,
                             role_arn: str, is_dir=False) -> bool:
    """
    download objects from s3 to local
    :param bucket: s3 bucket
    :param object_name: s3 object prefix
    :param file_path: prefix of where where objects to be saved, like '/tmp/'
    :param role_arn
    :return:
    """
    s3_client, _ = boto_client('s3', role_arn)
    config = TransferConfig(multipart_chunksize=(128 * 1024 * 1024), max_concurrency=3)
    transfer = S3Transfer(client=s3_client, config=config)
    contents = s3_client.list_objects(Bucket=bucket, Prefix=object_name)['Contents']
    if is_dir:
        zenyatta.common.util.make_dir(file_path + object_name)
    else:
        zenyatta.common.util.make_dir(file_path + '/'.join(object_name.split('/')[:-1]))
    try:
        for content in contents:
            if not content['Key'].endswith("/"):
                transfer.download_file(bucket, content['Key'], file_path + content['Key'])
                logging.info("downloading object from {bucket} {object_name} to "
                             "{file_path}{object_name}".format(**locals()))
        return True
    except ClientError as e:
        logging.info("failed to download objects from {bucket} {object_name} to "
                     "{file_path} {object_name}".format(**locals()) +
                     "{}".format(str(e)))
    return False


def remove_s3_object(bucket: str, object_name: str, role_arn: str) -> bool:
    """
    delete s3 object
    :param bucket: s3 bucket
    :param object_name: s3 object prefix
    :param role_arn
    :return:
    """
    s3_resource, _ = boto_resource('s3', role_arn)
    s3_client, _ = boto_client('s3', role_arn)
    try:
        contents = s3_client.list_objects(Bucket=bucket, Prefix=object_name)['Contents']
        for content in contents:
            response = s3_resource.Object(bucket, content['Key']).delete()
        logging.info("deleted object on {bucket} {object_name}".format(**locals()))
        return True
    except ClientError as e:
        logging.info("failed to delete objects from {bucket} {object_name}".format(**locals()) +
                     "{}".format(str(e)))
        return False
