#!python
from __future__ import print_function
import sys
import argparse
import base64
import getpass
import os
import shlex
import subprocess
import traceback
import boto3

from dockerutils import *
if sys.version_info < (3, 0):
    from ConfigParser import ConfigParser
else:
    from configparser import ConfigParser


def docker_login_aws(user, password, endpoint):
    return run('docker login -u {user} -p {password} {endpoint}'
               .format(user=user, password=password, endpoint=endpoint))


def docker_login_dockerhub():
    return run('docker login docker.com')


def docker_tag(image_name, image_tag, publication_tag):
    return run('docker tag {image_name}:{image_tag} {publication_tag}'
               .format(image_name=image_name, image_tag=image_tag, publication_tag=publication_tag))


def docker_push(publication_tag):
    return run('docker push {publication_tag}'.format(publication_tag=publication_tag))


def aws_create_repo_if_necessary(ecr_client, image_name):
    try:
        ecr_client.create_repository(repositoryName=image_name)
    except ecr_client.exceptions.RepositoryAlreadyExistsException:
        pass
    except Exception:
        traceback.print_exc()
        return -1
    return 0


def run(cmd):
    print('\n\n============================================================================')
    print('{cmd}\n\n'.format(cmd=cmd))
    return subprocess.call(shlex.split(cmd), cwd=os.getcwd())


if __name__ == '__main__':
    with cd(get_root_dir()):
        config = ConfigParser()
        config.optionxform = str
        config.read(os.path.join('docker', 'dockerutils.cfg'))

        image_types = get_image_types()

        parser = argparse.ArgumentParser()
        parser.add_argument("image", choices=image_types, help="image to publish")
        parser.add_argument("-i", "--image_name", help="use this image name rather than the default")
        args = parser.parse_args()

        image_name, image_tag = get_image_designation(args.image, config)
        if args.image_name:
            image_name = args.image_name

        push_config = {
            'image_repo': '',  # one of ['aws', 'docker_hub']
            'publication_tag': '',
        }
        if args.image in config.sections():
            for key in push_config.keys():
                if not push_config[key] and key in config.options(args.image):
                    push_config[key] = config.get(args.image, key)

        image_name, image_tag = get_image_designation(args.image, config)

        if not push_config['image_repo']:
            push_config['image_repo'] = 'docker'
        if not push_config['image_repo'] in ['aws', 'docker_hub']:
            print("Invalid publication target. Currently only supports 'aws' and 'docker_hub'")
            sys.exit(-1)

        if push_config['image_repo'] == 'aws':
            region = boto3.session.Session().region_name
            publication_tag = push_config['publication_tag'].format(
                account=boto3.client('sts').get_caller_identity()['Account'],
                region=region,
                image=image_name,
                tag=image_tag,
                user=getpass.getuser()
            )
            publication_tag = os.path.expandvars(publication_tag)
            ecr_client = boto3.client('ecr')
            auth_data = ecr_client.get_authorization_token()
            user, password = base64.b64decode(auth_data['authorizationData'][0]['authorizationToken']).decode().split(
                ':')
            endpoint = auth_data['authorizationData'][0]['proxyEndpoint']
            rc = docker_login_aws(user, password, endpoint)
            if rc == 0:
                aws_create_repo_if_necessary(ecr_client, image_name)
        else:
            publication_tag = push_config['publication_tag'].format(
                account='',
                region='',
                image=image_name,
                tag=image_tag,
                user=getpass.getuser()
            )
            publication_tag = os.path.expandvars(publication_tag)
            # TODO: should we pick up user and password for docker.com? Maybe via credstash?
            rc = docker_login_dockerhub()

        if rc == 0:
            rc = docker_tag(image_name, image_tag, publication_tag)
        if rc == 0:
            rc = docker_push(publication_tag)
    sys.exit(rc)
