#!/usr/bin/env python
import argparse
import yaml
import json
import sys
import re
import subprocess

GPU_DEVICE_PATTERN = re.compile(r'/dev/nvidia\d+')

# support Python 2 or 3
if sys.version_info[0] == 3:
    file_error = FileNotFoundError
else:
    file_error = IOError


def filehandle_if_exists_else_none(fname):
    try:
        return open(fname, 'r')
    except file_error:
        return None


def open_compose_file(fname):
    if not fname:
        return filehandle_if_exists_else_none('docker-compose.yaml') \
            or filehandle_if_exists_else_none('docker-compose.yml')
    else:
        return filehandle_if_exists_else_none(fname)


parser = argparse.ArgumentParser()
parser.add_argument('-f', '--file', metavar='INPUT_FILE', type=open_compose_file,
                    default='',
                    help='Specify an alternate input compose file')
parser.add_argument('-t', '--template', type=argparse.FileType('r'),
                    help='Specify Jinja2 template file from which compose file will be generated. '
                         '--template argument discards --file argument.')
parser.add_argument('-o', '--output', metavar='OUTPUT_FILE', type=argparse.FileType('w'),
                    default='jinja-compose.yml',
                    help='Specify an alternate output compose file (default: nvidia-docker-compose.yml)')
parser.add_argument('-G', '--generate', action='store_true',
                    help='Generate output compose file and exit, do not run docker-compose')

(args, extras) = parser.parse_known_args()

if args.file is None and args.template is None:
    print('Missing docker-compose file.')
    sys.exit(1)

# Rather hacky implementation to obtain the number of GPUs. Look for an improvement in the future
try:
    n_gpu = len(subprocess.check_output(['nvidia-smi', '-L'], timeout=10).decode().strip().split('\n'))
except (subprocess.TimeoutExpired, subprocess.CalledProcessError, file_error):
    # Error calling nvidia-smi. setting N_GPU to 0
    n_gpu = 0

print('Found {} GPUs'.format(n_gpu))


if args.template is not None:
    from jinja2 import Template
    content = Template(args.template.read()).render(N_GPU=n_gpu)
    config = yaml.load(content)
else:
    config = yaml.load(args.file)

if config is None:
    raise RuntimeError('Compose file is empty')

yaml.safe_dump(config, args.output, default_flow_style=False)

if not args.generate:
    from compose.cli.main import main as compose_main
    sys.argv[:] = ['docker-compose', '-f', args.output.name] + extras
    compose_main()