#!/usr/bin/env python
# -*- coding: utf-8 -*-

import getpass
import os
import time
import re
import sys

import click
import docker
from executor import execute
import yaml

CONFIG_PATH = os.path.join(os.path.expanduser("~"), ".kdk", "config.yaml")
CONFIG_SSH_KEY = os.path.join(os.path.expanduser("~"), ".kdk", "ssh", "id_rsa")

DOCKER_CLIENT = docker.from_env()

DEFAULT_CONFIG = {
    "image": {
        "repository": "ciscosso/kdk",
        "tag": "latest"
    },
    "docker": {
        "detach": True,
        "environment": {},
        "hostname": "kdk",
        "name": "kdk",
        "privileged": True,
        "stdin_open": True,
        "tty": True,
        "ports": {
            '2022': '2022',
        },
        "volumes": {},
    }
}


def load_config(config_path=CONFIG_PATH):
    return yaml.load(open(config_path).read())


class GroupOrderedHelp(click.Group):

    def __init__(self, *args, **kwargs):
        self.help_priorities = {}
        super(GroupOrderedHelp, self).__init__(*args, **kwargs)

    def get_help(self, ctx):
        self.list_commands = self.list_commands_for_help
        return super(GroupOrderedHelp, self).get_help(ctx)

    def list_commands_for_help(self, ctx):
        """reorder the list of commands when listing the help"""
        commands = super(GroupOrderedHelp, self).list_commands(ctx)
        return (c[1] for c in sorted((self.help_priorities.get(command, 1), command) for command in commands))

    def command(self, *args, **kwargs):
        """Behaves the same as `click.Group.command()` except capture a priority for listing command names in help.
        """
        help_priority = kwargs.pop("help_priority", 1)
        help_priorities = self.help_priorities

        def decorator(f):
            cmd = super(GroupOrderedHelp, self).command(*args, **kwargs)(f)
            help_priorities[cmd.name] = help_priority
            return cmd

        return decorator


@click.group(cls=GroupOrderedHelp)
@click.pass_context
def kdk(ctx):
    pass


@kdk.command(help_priority=1)
@click.pass_context
def init(ctx):
    """Init KDK - creates kdk config and pull latest image"""
    if not os.path.exists(os.path.dirname(CONFIG_PATH)):
        os.makedirs(os.path.dirname(CONFIG_PATH))

    if os.path.isfile(CONFIG_PATH):
        click.confirm(
            "A KDK config file already exists [{}].\nDo you want to continue and overwrite existing config?".format(
                CONFIG_PATH), abort=True)
    config = DEFAULT_CONFIG
    config["docker"]["environment"]["KDK_USERNAME"] = click.prompt("KDK username", default=getpass.getuser())
    config["image"]["repository"] = click.prompt("KDK image repository", default="ciscosso/kdk")
    config["image"]["tag"] = click.prompt("KDK image tag", default="debian-9.4-latest")
    config["docker"]["environment"]["KDK_DOTFILES_REPO"] = click.prompt("dotfiles_repo",
                                                                        default="https://github.com/cisco-sso/yadm-dotfiles.git")
    config["docker"]["environment"]["KDK_SHELL"] = click.prompt("shell", default="/bin/bash")

    # Volume mount the host ssh-pub-key to the guest authorized_keys
    config["docker"]["volumes"][os.path.dirname(CONFIG_SSH_KEY)+'/id_rsa.pub'] = {
        'bind': "/home/{}/.ssh/authorized_keys".format(
            config["docker"]["environment"]["KDK_USERNAME"]),
        'mode': 'rw'
    }

    # Write the config to file
    yaml.dump(config, open(CONFIG_PATH, "w"), default_flow_style=False)
    click.echo("KDK config saved to {}".format(CONFIG_PATH))

    # Pull the image
    ctx.invoke(pull)

    # Ensure an ssh key has been generated
    if os.path.isfile(CONFIG_SSH_KEY) is False:
        os.makedirs(os.path.dirname(CONFIG_SSH_KEY), exist_ok=True)
        cmd='docker run --rm -v {}:/tmp/ssh {}:{} bash -c \'echo -e "\\n\\n\\n" | /usr/bin/ssh-keygen -b 4096 -N "" -f /tmp/ssh/id_rsa\''.format(
            os.path.dirname(CONFIG_SSH_KEY),
            config["image"]["repository"],
            config["image"]["tag"],
        )
        try:
            execute(cmd)
        except:
            pass

    click.echo("KDK init complete.")


@kdk.command(help_priority=2)
@click.pass_context
def pull(ctx):
    """Pull the latest/configured KDK image"""
    config = load_config()
    repository = config["image"]["repository"]
    tag = config["image"]["tag"]
    click.echo("Pulling KDK image [{}:{}]. This may take a moment...".format(repository, tag))
    DOCKER_CLIENT.images.pull(repository=repository, tag=tag)
    click.echo("Finished pulling KDK image {}:{}".format(repository, tag))


@kdk.command(help_priority=3)
@click.pass_context
def up(ctx, interactive=True, repository=None, tag=None):
    """Start a KDK container"""
    kdk_container = None
    try:
        kdk_container = DOCKER_CLIENT.containers.get("kdk")
    except:
        pass
    if kdk_container:
        click.echo("KDK container is already running.")
        exit(1)
    config = load_config()
    if repository:
        config["image"]["repository"] = repository
    if tag:
        config["image"]["tag"] = tag
    if interactive:
        user_kdk_image_pattern = re.compile("{}-kdk:\d+".format(getpass.getuser()))
        images = DOCKER_CLIENT.images.list()
        user_kdk_images = []
        user_kdk_image_tags = []
        for i in images:
            if "RepoTags" in i.attrs and i.attrs["RepoTags"]:
                for image_tag in i.attrs["RepoTags"]:
                    if user_kdk_image_pattern.match(image_tag):
                        user_kdk_images.append(i)
                        user_kdk_image_tags.append(image_tag)
                        break
        if user_kdk_images:
            click.echo("KDK user snapshots exist\n{}:".format(user_kdk_image_tags))
            desired_image = click.prompt("Specify KDK image.",
                                         default="{}:{}".format(config["image"]["repository"], config["image"]["tag"]))
            config["image"]["repository"], config["image"]["tag"] = desired_image.split(":")
    click.echo("Starting KDK container with following config:\n{}".format(yaml.dump(config, default_flow_style=False)))
    DOCKER_CLIENT.containers.run(
        image="{}:{}".format(config["image"]["repository"], config["image"]["tag"]),
        **config["docker"])
    ctx.invoke(provision)


@kdk.command(help_priority=4)
@click.pass_context
def provision(ctx):
    """Provision KDK user"""
    click.echo("Provisioning KDK user (this may take a moment)...")
    try:
        execute("docker exec -it kdk /usr/local/bin/provision-user")
    except:
        click.echo("ERROR: Failed to provision KDK user")
        raise


@kdk.command(help_priority=5)
@click.pass_context
def ssh(ctx):
    """SSH to running KDK container"""
    config = load_config()
    user = config["docker"]["environment"]["KDK_USERNAME"]

    cmd = "ssh {user}@localhost -A -p 2022 -i ~/.kdk/ssh/id_rsa -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ".format(**locals())
    try:
        click.echo(cmd)
        execute(cmd)
    except:
        pass


@kdk.command(help_priority=6)
@click.pass_context
def snapshot(ctx, interactive=True):
    """Create a snapshot of running KDK container"""
    repository = "{}-kdk".format(getpass.getuser())
    tag = int(time.time())
    try:
        kdk_container = DOCKER_CLIENT.containers.get("kdk")
    except:
        if interactive:
            click.echo("ERROR: Failed to snapshot KDK. Confirm the container is running.")
            exit(1)
        else:
            raise
    kdk_container.commit(repository=repository, tag=tag)
    click.echo("KDK snapshot created [{}:{}]".format(repository, tag))
    return {"repository": repository, "tag": tag}


@kdk.command(help_priority=7)
@click.pass_context
def prune(ctx):
    """Prune unused KDK images"""
    running_kdk_container = None
    try:
        running_kdk_container = DOCKER_CLIENT.containers.get("kdk")
        running_kdk_container = running_kdk_container.attrs["Config"]["Image"]
    except:
        pass
    user_kdk_image_pattern = re.compile("{}-kdk:\d+".format(getpass.getuser()))
    user_kdk_image_tags = []
    images = DOCKER_CLIENT.images.list()
    for i in images:
        if "RepoTags" in i.attrs and i.attrs["RepoTags"]:
            for image_tag in i.attrs["RepoTags"]:
                if user_kdk_image_pattern.match(image_tag):
                    user_kdk_image_tags.append(image_tag)
                    break
    if running_kdk_container:
        stale_user_kdk_image_tags = [i for i in user_kdk_image_tags if i != running_kdk_container]
    else:
        stale_user_kdk_image_tags = [i for i in user_kdk_image_tags]
    if stale_user_kdk_image_tags:
        click.confirm("Prune unused KDK images?\n{}".format(stale_user_kdk_image_tags), abort=True)
        for tag in stale_user_kdk_image_tags:
            DOCKER_CLIENT.images.remove(tag)
    else:
        click.echo("No stale images found.")


@kdk.command(help_priority=8)
@click.pass_context
def destroy(ctx, interactive=True):
    """Destroy the running KDK container"""
    if interactive:
        click.confirm("Destroy the running KDK container and associated image?", abort=True)
    try:
        kdk_container = DOCKER_CLIENT.containers.get("kdk")
    except:
        click.echo("ERROR: Failed to destroy KDK. Confirm the container is running.")
        exit(1)
    kdk_container.remove(force=True)
    click.echo("Destroyed KDK container and associated image.")


@kdk.command(help_priority=9)
@click.pass_context
def restart(ctx):
    """Restart the running KDK container (useful for config changes)"""
    try:
        _snapshot = ctx.invoke(snapshot, interactive=False)
    except:
        click.echo("ERROR: Failed to restart -- unable to snapshot running KDK container")
        exit(1)
    ctx.invoke(destroy, interactive=False)
    ctx.invoke(up, interactive=False, **_snapshot)
    click.echo("KDK container restarted.")
    ctx.invoke(prune)


if __name__ == "__main__":
    kdk()
