#!python

# Python script to run containers associated with a project
#
from __future__ import print_function
import sys
import argparse
import datetime
import getpass
import os
import shlex
import subprocess

from string import Template

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


_base_cmd = 'docker run {init} --name {name} {environment} {keep_container} {interactive} {gpu} {network} ' \
            '{volumes} {ports} {args} {image_name}:{image_tag}'

def fetch_env_variables(config, image, args_env=None):
    # a special use case. retrieve env variables from a config server.
    env_vars = {}
    if args_env:
        for arg_env in args_env:
            env_var = arg_env.split('=')
            env_vars[env_var[0]] = env_var[1]
    if config.has_section(image) and config.has_option(image, 'env'):
        env_section = config.get(image, 'env')
        for evars in config.options(env_section):
            env_vars[evars] = config.get(env_section, evars)
    return ' '.join(['-e {key}={value}'.format(key=key, value=value) for key, value in env_vars.items()])


def populate_aws_env_variables():
    AWS_ENVS = [
        'AWS_ACCESS_KEY_ID',
        'AWS_SECRET_ACCESS_KEY',
        'AWS_DEFAULT_REGION',
        'AWS_REGION',
        'AWS_SESSION_TOKEN',
        'AWS_SECURITY_TOKEN'
    ]

    variables = {}
    for test_env in AWS_ENVS:
        if os.getenv(test_env):
            variables[test_env] = "${" + test_env + "}"

    return ' '.join(['-e {key}={value}'.format(key=k, value=v) for k,v in variables.items()])


def run(mode, image_name, image_tag, **kwargs):
    user = getpass.getuser()
    volumes = kwargs['volumes'].format(
        project_root=get_root_dir(),
        user=user,
        project=os.path.split(get_root_dir())[1]
    )
    volumes = os.path.expandvars(volumes)

    if kwargs['network']:
        kwargs['network'] = "--network {network}".format(network=kwargs['network'])

    timestamp = datetime.datetime.now().strftime("%y-%m-%d_%H.%M.%S")
    cmd = _base_cmd.format(image_name=image_name,
                           image_tag=image_tag,
                           name="{user}_{mode}_{timestamp}".format(
                               user=getpass.getuser(), mode=mode, timestamp=timestamp),
                           keep_container=kwargs['keep_container'],
                           interactive=kwargs['interactive'],
                           environment=kwargs['environment'],
                           network=kwargs['network'],
                           ports=kwargs['ports'],
                           args=kwargs['args'],
                           volumes=volumes,
                           gpu=kwargs['gpu'],
                           init=kwargs['init'])
    print('\n\n============================================================================')
    print("{cmd} {kwcmd}\n\n".format(cmd=cmd, kwcmd=kwargs['cmd'] or ''))

    # Since we are using secure env values I don't want those to print in the above commmand, but
    #   they need to be expanded for the subprocess.call
    expanded_cmd = Template(cmd).substitute(os.environ)
    expanded_kwargs_cmd = Template(kwargs['cmd'] or '').substitute(os.environ)

    all_args = shlex.split(expanded_cmd) + shlex.split(expanded_kwargs_cmd)
    return subprocess.call(all_args, cwd=os.getcwd())


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

        if 'run_image' in config.sections() and 'synthetic_images' in config.options('run_image'):
            image_types = get_image_types(config.get('run_image', 'synthetic_images').split(','))
        else:
            image_types = get_image_types()

        parser = argparse.ArgumentParser()
        parser.add_argument("image", choices=image_types, help="Docker image to run")
        parser.add_argument("-a", "--args", help="general docker arguments", default='')
        parser.add_argument("-k", "--keep", help="keep the image after execution", action='store_true')
        parser.add_argument("-c", "--command", help="Command for image override")
        parser.add_argument("-n", "--network", help="Network for image override", default='')
        parser.add_argument("-g", "--use-gpu", dest='use_gpu', default=False, action='store_true',
                            help="Start the container with gpu support")
        parser.add_argument("-e", "--env", action='append',
                            help="environment variables to pass to running container, e.g. foo=bar")
        args = parser.parse_args()

        if args.use_gpu:
            gpu = '--runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all'
        else:
            gpu = ''

        run_config = {
            'environment': fetch_env_variables(config, args.image, args.env) + ' ' + populate_aws_env_variables(),
            'keep_container': args.keep or '--rm',
            'interactive': '-d' if args.keep else '-it',
            'gpu': gpu,
            'network': args.network or '',
            'volumes': '',
            'ports': '',
            'cmd': args.command or None,
            'init': '--init',
            'args': args.args
        }

        is_docked = bool(os.environ.get('DOCKER_IP'))

        if args.image in config.sections():
            for key in run_config.keys():
                docked_key = "{key}_docked".format(key=key)
                if is_docked and not run_config[key] and docked_key in config.options(args.image):
                    run_config[key] = config.get(args.image, docked_key)
                elif not run_config[key] and key in config.options(args.image):
                    run_config[key] = config.get(args.image, key)

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

        sys.exit(run(args.image, image_name, image_tag, **run_config))
