import argparse
import glob
import pathlib
import boto3
from io import BytesIO

S3 = boto3.client('s3')

DO_NOT_COPY = ['airflow.db']


def push_config(bucket: str, key: str, config_dir: str):

    for conf_file in glob.glob(config_dir + "/*"):
        file_name = pathlib.Path(conf_file).name
        if file_name in DO_NOT_COPY:
            print("skipping: {file_name}".format(**locals()))
            continue
        output_key = key + "/" + file_name
        try:
            S3.put_object(Body=open(conf_file, 'r').read(), Bucket=bucket, Key=output_key)
            print("uploaded " + conf_file + " successfully")
        except Exception as s3_ex:
            print("couldn't put file: {conf_file} in s3://{bucket}/{output_key}".format(**locals()))
            print(s3_ex)
            raise s3_ex


def pull_config(bucket: str, key: str, config_dir: str):

    objects = S3.list_objects_v2(Bucket=bucket, Prefix=key)
    for item in objects['Contents']:
        response = S3.get_object(Bucket=bucket, Key=item['Key'])
        file_contents = BytesIO(response['Body'].read()).read().decode()
        output_path = pathlib.Path(config_dir + "/" + pathlib.Path(item['Key']).name)
        print("writing {} to {}".format(item["Key"], str(output_path)))
        with open(str(output_path), 'w') as f:
            f.write(file_contents)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="push all of the airflow config to s3 or pull all of it "
                                                 "to /etc/zenyatta")
    parser.add_argument('action', choices=['push', 'pull'])
    parser.add_argument('--bucket', required=False, dest='bucket', action='store', help='s3 bucket')
    parser.add_argument('--key', required=True, dest='key', action='store', help='s3 key')
    parser.add_argument('--config-dir', required=False, dest='config_dir', default='/etc/zenyatta',
                        action='store', help='directory where config is')
    args = parser.parse_args()

    if not args.bucket:
        from zenyatta.common import get_connections
        args.bucket = get_connections()['aws']['s3-input']['bucket']

    if 'push' in args.action:
        push_config(args.bucket, args.key, args.config_dir)
    elif 'pull' in args.action:
        pull_config(args.bucket, args.key, args.config_dir)
