#!/usr/bin/env python
# -*- config:utf-8 -*-
"""
$ ./add_s3_tag.py --help
usage: add_s3_tag.py [-h] [--profile PROFILE] [--tag-value TAG_VALUE]
                     [--tag-key TAG_KEY] [--query QUERY] [--dry-run]

optional arguments:
  -h, --help            show this help message and exit
  --profile PROFILE
  --tag-value TAG_VALUE
  --tag-key TAG_KEY
  --query QUERY
  --dry-run
"""
import argparse
import copy
import boto3
from boto3.session import Session
from botocore.exceptions import ClientError


def get_args():
    """ コマンド引数のパースを行う
    """
    psr = argparse.ArgumentParser()
    psr.add_argument('--profile', default='default')
    psr.add_argument('--tag-value', default='')
    psr.add_argument('--tag-key', default='')
    psr.add_argument('--query', default='')
    psr.add_argument('--dry-run', action='store_true')
    return psr.parse_args()


def dict_list_uniquify(target_list):
    """ 辞書の配列から重複要素を排除する
    """

    list_tuppled = [tuple(d.items()) for d in target_list]
    list_tuppled_unique = list(set(list_tuppled))
    list_unique = [dict(t) for t in list_tuppled_unique]

    return list_unique


def add_tag(s3, bucket_name, tag_key, tag_value, dry_run):
    """ バケットにタグ付けする
    """

    print(bucket_name)

    try:
        bucket_tagging = s3.BucketTagging(bucket_name)
        tags = bucket_tagging.tag_set
        tags.append({'Key': tag_key, 'Value': tag_value})
        add_tags = dict_list_uniquify(tags)
    except ClientError as e:
        if e.response['Error']['Code'] == 'NoSuchTagSet':
            add_tags = [{'Key': tag_key, 'Value': tag_value}]
        else:
            print("Unexpected error: %s" % e)

    print(add_tags)
    if not dry_run:
        print('Add tags...')
        bucket_tagging.put(Tagging={'TagSet': add_tags})


def main():
    args = get_args()

    session = Session(profile_name=args.profile)

    s3 = session.client('s3')
    s3_resource = session.resource('s3')
    resp = s3.list_buckets()
    bucket_names = [
        bucket['Name'] for bucket in resp['Buckets']
        if bucket['Name'].find(args.query) >= 0
    ]
    _ = list(
        map(
            lambda b: add_tag(s3_resource, b, args.tag_key, args.tag_value, args.dry_run),
            bucket_names))


if __name__ == '__main__': main()