#!/usr/bin/env python3.6
import argparse
import io
import logging
import os
import re
import sh
import shlex
import sys
import tempfile

import yaml

from distutils.spawn import find_executable

# regular expression for finding variables in docker compose files
VAR_RE = re.compile(r'\${(?P<varname>.*)}')

COMPOSE_FILES = ['docker-compose-stack.yml', 'docker-compose.yml']
DC_CONFIG_FILES = ['dc.yml', 'dc-overlay.yml']

# directory to keep all compose files in
COMPOSE_DIR = 'compose'

PROJECT_NAME = None

# check to see if an overlay file is provided in the environment
DC_CONFIG_FILE = os.environ.get('DC_CONFIG_FILE', '')
if not DC_CONFIG_FILE:
    # find one
    for _file in DC_CONFIG_FILES:
        for _dir in ('.', COMPOSE_DIR):
            compose_path = f'{_dir}/{_file}'
            if os.path.exists(compose_path):
                DC_CONFIG_FILE = _file

# find where the compose file is
for _file in COMPOSE_FILES:
    for _dir in ('.', COMPOSE_DIR):
        compose_path = f'{_dir}/{_file}'
        if os.path.exists(compose_path):
            PROJECT_NAME = os.path.basename(os.getcwd())
            # print(f'PROJECT_NAME={PROJECT_NAME}')

            os.chdir(_dir)

ENVIRONMENT_ROOT = os.path.expanduser('~/.docker/_environments')
DC_ENVIRONMENT = os.environ.get('DC_ENVIRONMENT', ENVIRONMENT_ROOT)


class CommandError(Exception):
    pass


def check_environment(args):
    logger = logging.getLogger('main')

    if not args.environment:
        return

    environment_path = get_environment_path(args)

    # the stack environment includes the environment name
    stack_env = {
        'DOCKER_STACK': args.environment or PROJECT_NAME,
    }

    if os.path.exists(environment_path):
        with open(environment_path) as fh:
            for line in fh.read().splitlines():
                if line.startswith('#'):
                    continue

                try:
                    key, value = line.split('=', 1)
                except ValueError as exc:
                    print(f'unable to split line={line}')
                    raise
                stack_env[key] = value
    else:
        raise CommandError(f'environment={args.environment} in {DC_ENVIRONMENT} not found')

    if args.tag_version:
        tag_version_command = getattr(sh, 'tag-version')
        tag_version = tag_version_command().stdout.decode('utf8').strip()

        stack_env['VERSION'] = tag_version

    if args.tag_docker_image:
        if not args.tag_version:
            raise CommandError('cannot tag docker image without setting --tag-version')

        docker_image = os.environ.get('DOCKER_IMAGE', stack_env['DOCKER_IMAGE']).rsplit(':', 1)[0]
        tagged_docker_image = f'{docker_image}:{tag_version}'
        stack_env['DOCKER_IMAGE'] = tagged_docker_image

        logger.info(f'tagged_docker_image={tagged_docker_image}')

    os.environ.update(stack_env)

    # for k, v in stack_env.items():
    #     print(f'{k}={v}')

    return stack_env


def deploy(args, filenames):
    if len(filenames) > 1:
        return 'deploy only supports a single compose file'

    compose_file = filenames[0]
    if not os.path.exists(compose_file):
        return f'compose_file={compose_file} not found'

    with open(compose_file) as fh:
        content = fh.read()

        previous_idx = 0
        rendered = ''
        for x in VAR_RE.finditer(content):
            rendered += content[previous_idx:x.start('varname')-2]  # -2 to get rid of variable's `${`

            varname = x.group('varname')
            try:
                rendered += os.environ[varname]
            except KeyError:
                sys.exit(f'varname={varname} not in environment')

            previous_idx = x.end('varname') + 1  # +1 to get rid of variable's `}`

        rendered += content[previous_idx:]

    fh = tempfile.NamedTemporaryFile()
    fh.write(rendered.encode('utf8'))
    fh.flush()

    stack_name = args.environment or PROJECT_NAME

    command = f"""docker stack deploy
      --prune
      --with-registry-auth
      --compose-file {fh.name}
      {stack_name}"""
    command_split = shlex.split(command)

    print(command)

    if args.noop:
        print(rendered)
    else:
        executable = getattr(sh, command_split[0])
        executable(*command_split[1:], _env=os.environ)


def get_environment_path(args):
    return os.path.join(DC_ENVIRONMENT, args.environment)


def get_overlay_filenames(overlay):
    logger = logging.getLogger('get_overlay_filenames')

    overlay_filenames = []

    applied = []
    for item in overlay:
        if item in applied:
            continue

        applied.append(item)

        path = None
        if isinstance(item, dict):
            name = item['name']
            path = item['path']
        else:
            name = item

        # join path and name if the path is given
        if path:
            _filename = os.path.join(path, name)
        else:
            _filename = name

        if _filename and os.path.exists(_filename) and os.path.isfile(_filename):
            overlay_filenames.append(_filename)
        else:
            # prefix partial with a dot in order to complete the name
            if name:
                name = '.{}'.format(name)
            else:
                name = ''

            _filename = 'docker-compose{}.yml'.format(name)
            if path:
                _filename = os.path.join(path, _filename)

            logging.debug('_filename={}'.format(_filename))

            if os.path.exists(_filename):
                overlay_filenames.append(_filename)
            else:
                logger.warning(f'filename={_filename} does not exist, skipping')

    # check to see if any filenames were found, else default to docker-compose.yml
    if not overlay_filenames:
        overlay_filenames.append('docker-compose.yml')

    return overlay_filenames


def main(args, docker_args):
    logger = logging.getLogger('main')

    # make sure just a task or profile is given
    if all([args.profile, args.task]) or not any([args.profile, args.task]):
        args.profile = 'default'
        logger.warning('using default profile')

    data = None
    overlay = []

    check_environment(args)

    if os.path.exists(DC_CONFIG_FILE):
        with open(DC_CONFIG_FILE, 'r') as fh:
            data = yaml.load(fh)

    if data:
        if args.task:
            task = data['tasks'][args.task]
            command = shlex.split(task['command'])

            return init(command=command + args.docker_args)

        try:
            overlay_data = data.get('profiles') or data['overlay']
        except KeyError:
            return 'ERROR: define profiles section in dc.yml'

        # check which profile to use
        profile = args.profile
        try:
            overlay.extend(overlay_data[profile])
        except KeyError:
            print('profile {} not found'.format(profile))
            return -1
        except TypeError:
            print('place overlays within named profiles under the `overlay` section')
            return -1
    elif args.task:
        return 'ERROR: cannot run task without overlay file'

    for item in (args.add_overlay or []):
        overlay.extend([x.strip() for x in item.split(',')])

    logger.debug('overlay={}'.format(overlay))

    filenames = get_overlay_filenames(overlay)

    if args.deploy:
        return deploy(args, filenames)
    else:
        return run_compose(args, filenames, docker_args)

def run_compose(args, filenames, docker_args):
    command = ['docker-compose']
    command[0] = find_executable(command[0])

    if args.project_name:
        command.extend(['--project-name', args.project_name])

    for _filename in filenames:
        command.extend(['-f', _filename])

    command.extend(docker_args)

    print(' '.join(command), file=sys.stderr)

    if not args.noop:
        # os.execve(command[0], command, os.environ)
        proc = getattr(sh, command[0])
        proc(*command[1:], _env=os.environ, _fg=True)

        if args.write_tag:
            write_tag(args)

        if args.push:
            sh.docker('push', os.environ['DOCKER_IMAGE'], _fg=True)


def write_tag(args):
    buf = io.StringIO()

    env_file = get_environment_path(args)
    with open(env_file) as fh:
        while True:
            line = fh.readline()
            if line == '':
                break

            line = line.strip()

            if line.startswith('DOCKER_IMAGE='):
                line = f'DOCKER_IMAGE={os.environ["DOCKER_IMAGE"]}'

            buf.write(f'{line}\n')

    with open(env_file, 'w') as fh:
        fh.write(buf.getvalue())


def init(command=None):
    parser = argparse.ArgumentParser()

    # required
    parser.add_argument('-p', '--profile', help='profile to use in overlay')

    # optional
    parser.add_argument('-e', '--environment', help='load up environment prior to running command')
    parser.add_argument('-a', '--add-overlay', help='additional file to overlay', action='append')
    parser.add_argument('-n', '--noop', action='store_true', help='just print command, do not execute')
    parser.add_argument('--project-name', default=PROJECT_NAME, help=f'the projet name to use, default={PROJECT_NAME}')
    parser.add_argument('--task', help='run a task defined in dc-overlay.yml')
    parser.add_argument('--deploy', action='store_true', help='execute a deployment')
    parser.add_argument(
        '--tag-version',
        action='store_true',
        help='set VERSION environment var with tag-version output'
    )
    parser.add_argument(
        '--tag-docker-image',
        action='store_true',
        help='when --tag-version is enabled, replace the tag in DOCKER_IMAGE environment variable'
    )
    parser.add_argument(
        '--write-tag',
        action='store_true',
        help='when --tag-docker-image is enabled, write DOCKER_IMAGE to the environment file'
    )
    parser.add_argument(
        '--push',
        action='store_true',
        help='when enabled, push DOCKER_IMAGE'
    )

    parser.add_argument('docker_args', nargs=argparse.REMAINDER)

    if command is None:
        command = sys.argv

    args = parser.parse_args(command[1:])

    # hack when no profile or task is given
    # take the first docker_args argument as the profile
    if not any([args.profile, args.task]) and len(args.docker_args) > 1:
        item = args.docker_args.pop(0)
        if item == 'task':
            args.task = args.docker_args.pop(0)
        else:
            args.profile = item

    # hack when there is no `dc.yml`
    # this is just so that the dc command can be used even when there is no overlay file in the project
    if not os.path.exists(DC_CONFIG_FILE):
        if args.profile:
            args.docker_args.insert(0, args.profile)

        args.profile = '-'

    return main(args, args.docker_args)


if __name__ == '__main__':
    log_level=os.environ.get('LOG_LEVEL', 'info')

    logging.basicConfig(level=getattr(logging, log_level.upper()))

    sys.exit(init())
