#!/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

from boltons.iterutils import remap, get_path, default_enter, default_visit

SCRIPT_DIR, SCRIPT_NAME = os.path.split(__file__)
MERGED_YML = 'docker-compose-merged.yml'

# 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
_found = False
for _file in COMPOSE_FILES:
    for _dir in ('.', COMPOSE_DIR):
        compose_path = f'{_dir}/{_file}'
        if os.path.exists(compose_path):
            _found = True

            PROJECT_NAME = os.path.basename(os.getcwd())
            # print(f'PROJECT_NAME={PROJECT_NAME}')

            os.chdir(_dir)

            break
    if _found:
        break

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('check_environment')

    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():
                # skip empty lines
                if line.strip() == '':
                    continue

                # skip commented lines
                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):
    logger = logging.getLogger('deploy')

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

    stack_name = args.environment or PROJECT_NAME

    command = f"""docker stack deploy
      --prune
      --with-registry-auth
      --compose-file {filenames[0]}
      {stack_name}"""
    command_split = shlex.split(command)

    logger.info(command)

    if not args.dry_run:
        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, prev_args=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)

    # merge multiple files together so that deploying stacks works
    # https://github.com/moby/moby/issues/30127
    if len(filenames) > 1:
        yaml_contents = []

        for item in filenames:
            with open(item, 'r') as fh:
                yaml_contents.append(yaml.load(fh))

        merged = remerge(yaml_contents)
        content = yaml.dump(merged, default_flow_style=False)
    else:
        with open(filenames[0], 'r') as fh:
            content = fh.read()

    # render the file
    rendered = render(content)

    fh = open(MERGED_YML, 'w')
    fh.write(rendered)
    fh.flush()

    # when executing a dry run, list out the environment and the merged docker-compose file
    if args.dry_run:
        print(f'environment:\n----------\n')
        print('\n'.join(sorted([f'{k}={v}' for k, v in os.environ.items()])))
        print(f'\n----------\n')

        print(f'{MERGED_YML}:\n----------\n')
        with open(MERGED_YML, 'r') as fh:
            print(fh.read())
        print(f'\n----------\n')

    try:
        if args.deploy:
            return deploy(args, [fh.name])
        else:
            return run_compose(args, [fh.name], docker_args)
    finally:
        os.remove(MERGED_YML)


# https://gist.github.com/mahmoud/db02d16ac89fa401b968
def remerge(target_list, sourced=False):
    """Takes a list of containers (e.g., dicts) and merges them using
    boltons.iterutils.remap. Containers later in the list take
    precedence (last-wins).
    By default, returns a new, merged top-level container. With the
    *sourced* option, `remerge` expects a list of (*name*, container*)
    pairs, and will return a source map: a dictionary mapping between
    path and the name of the container it came from.
    """

    if not sourced:
        target_list = [(id(t), t) for t in target_list]

    ret = None
    source_map = {}

    def remerge_enter(path, key, value):
        new_parent, new_items = default_enter(path, key, value)
        if ret and not path and key is None:
            new_parent = ret
        try:
            cur_val = get_path(ret, path + (key,))
        except KeyError:
            pass
        else:
            # TODO: type check?
            new_parent = cur_val

        if isinstance(value, list):
            # lists are purely additive. See https://github.com/mahmoud/boltons/issues/81
            new_parent.extend(value)
            new_items = []

        return new_parent, new_items

    for t_name, target in target_list:
        if sourced:
            def remerge_visit(path, key, value):
                source_map[path + (key,)] = t_name
                return True
        else:
            remerge_visit = default_visit

        ret = remap(target, enter=remerge_enter, visit=remerge_visit)

    if not sourced:
        return ret
    return ret, source_map


def render(content):
    """
    Renders the variables in the file
    """
    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:]

    return rendered


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

    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)

    logger.info(' '.join(command))

    if not args.dry_run:
        try:
            # os.execve(command[0], command, os.environ)
            proc = getattr(sh, command[0])
            proc(*command[1:], _env=os.environ, _fg=True)
        except:
            print(command)
            print(sh.cat(MERGED_YML).stdout.decode('utf8'))

            raise

    if args.write_tag:
        write_tag(args)

    if args.push:
        docker_image = os.environ['DOCKER_IMAGE']

        if args.dry_run:
            logger.info(f'docker push {docker_image}')
        else:
            sh.docker('push', docker_image, _fg=True)


def write_tag(args):
    logger = logging.getLogger('write_tag')
    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='):
                new_line = f'DOCKER_IMAGE={os.environ["DOCKER_IMAGE"]}'

                logger.info(f'replacing line={line} with new_line={new_line}')

                line = new_line

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

    if not args.dry_run:
        with open(env_file, 'w') as fh:
            fh.write(buf.getvalue())


def init(command=None, prev_args=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', '--dry-run',
        action='store_true', dest='dry_run',
        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

    if not command[0].endswith(SCRIPT_NAME):
        return run_external_command(command)

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

    # if previous arguments were given, check them for certain args
    if prev_args:
        args.dry_run = prev_args.dry_run

    # 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)


def run_external_command(command):
    return getattr(sh, command[0])(command[1:], _fg=True)


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

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

    sys.exit(init())
