#!python

# Python script to run containers associated with a project
#
import argparse
import configparser
import datetime
import getpass
import json
import os
import shlex
import subprocess
import sys
import urllib.parse
import urllib.request

from dockerutils import *

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

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 'env_var_names' in config.sections():
        endpoint = f"http://{os.environ.get(config['env_var_names']['remote_ip'])}:" \
                   f"{os.environ.get(config['env_var_names']['remote_port'], '5000')}/config"
        try:
            with urllib.request.urlopen(endpoint, timeout=1) as response_f:
                response = json.load(response_f)
                env_vars = response.get('env_variables') or {}
        except urllib.error.URLError:
            pass
    if image in config and 'env' in config[image]:
        env_section = config[image]['env']
        for vars in config[env_section]:
            env_vars[vars] = config[env_section][vars]
    return ' '.join([f'-e {key}={value}' for key, value in env_vars.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'] = f"--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=f"{getpass.getuser()}_{mode}_{timestamp}",
                           keep_container=kwargs['keep_container'],
                           interactive=kwargs['interactive'],
                           environment=kwargs['environment'],
                           network=kwargs['network'],
                           ports=kwargs['ports'],
                           volumes=volumes,
                           gpu=kwargs['gpu'],
                           cmd=kwargs['cmd'],
                           init=kwargs['init'])
    print('\n\n============================================================================')
    print(f'{cmd}\n\n')
    return subprocess.call(shlex.split(cmd), cwd=os.getcwd())


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

        if 'run_image' in config and 'synthetic_images' in config['run_image']:
            image_types = get_image_types(config['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("-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("-p", "--pre", default=False, action='store_true',
                            help="use pre 1.25 API")
        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 -e CPU_GPU_ENV=/gpu-env'
        else:
            gpu = ''
        if args.pre:
            init = ''
        else:
            init = '--init'

        run_config = {
            'environment': fetch_env_variables(config, args.image, args.env),
            '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 '',
            'init': init
        }

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

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

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

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