#!/usr/bin/python -Es
# Copyright (C) 2014-2015 Red Hat
# AUTHOR: Dan Walsh <dwalsh@redhat.com>
# see file 'COPYING' for use and warranty information
#
# atomic is a tool for managing Atomic Systems and Containers
#
#    This program is free software; you can redistribute it and/or
#    modify it under the terms of the GNU General Public License as
#    published by the Free Software Foundation; either version 2 of
#    the License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
#                                        02111-1307  USA
#
#
import sys, os
import argparse
import gettext
import docker
import json
import subprocess

PROGNAME="atomic"
gettext.bindtextdomain(PROGNAME, "/usr/share/locale")
gettext.textdomain(PROGNAME)
try:
    gettext.install(PROGNAME,
                    unicode=True,
                    codeset = 'utf-8')
except TypeError:
    # Failover to python3 install
    gettext.install(PROGNAME,
                    codeset = 'utf-8')
except IOError:
    import builtins
    builtins.__dict__['_'] = str

class Atomic(object):
    INSTALL_ARGS = ["/usr/bin/docker", "run",
                    "-t",
                    "-i",
                    "--rm",
                    "--privileged",
                    "-v", "/:/host",
                    "--net=host",
                    "--ipc=host",
                    "--pid=host",
                    "-e", "HOST=/host",
                    "-e", "NAME=NAME",
                    "-e", "IMAGE=IMAGE",
                    "-v", "${CONFDIR}:/etc/NAME",
                    "-v", "${LOGDIR}:/var/log/NAME",
                    "-v", "${DATADIR}:/var/lib/NAME",
                    "-e", "CONFDIR=${CONFDIR}",
                    "-e", "LOGDIR=${LOGDIR}",
                    "-e", "DATADIR=${DATADIR}",
                    "--name", "NAME",
                    "IMAGE"]

    SPC_ARGS = ["/usr/bin/docker", "run",
                "-t",
                "-i",
                "--rm",
                "--privileged",
                "-v", "/:/host",
                "-v", "/run:/run",
                "-v", "/etc/localtime:/etc/localtime",
                "--net=host",
                "--ipc=host",
                "--pid=host",
                "-e", "HOST=/host",
                "-e", "NAME=NAME",
                "-e", "IMAGE=IMAGE",
                "IMAGE" ]

    RUN_ARGS = ["/usr/bin/docker", "create",
                "-t",
                "-i",
                "--name", "NAME",
                "IMAGE" ]

    def __init__(self):
        self.d = docker.Client()
        self.name = None
        self.image = None
        self.spc = False
        self.inspect = None

    def writeOut(self, output):
        sys.stdout.flush()
        sys.stdout.write(output + "\n")

    def update(self):
        return subprocess.check_call(["/usr/bin/docker", "pull", self.image])

    def pull(self):
        prevstatus = ""
        prev = ""
        for line in self.d.pull(self.image, stream=True):
            bar = json.loads(line)
            status = bar['status']
            if prevstatus != status:
                print status
            if 'id' not in bar:
                continue
            if status == "Downloading":
                self.writeOut(bar['progress'] + " ")
            elif status == "Extracting":
                self.writeOut("Extracting: " + bar['id'])
            elif status == "Pull complete":
                pass
            elif status.startswith("Pulling"):
                self.writeOut("Pulling: " + bar['id'])

            prevstatus = status
        print("")

    def set_args(self, args):
        self.args=args
        try:
            self.image = args.image
        except:
            pass
        try:
            self.command = args.command
        except:
            self.command = None

        try:
            self.spc = args.spc
        except:
            self.spc = False

        try:
            self.name = args.name
        except:
            pass

        if not self.name and self.image is not None:
            self.name = self.image.split("/")[-1].split(":")[0]
            if self.spc:
                self.name = self.name + "-spc"

    
    def _getconfig(self, key, default=None):
        assert self.inspect is not None
        cfg = self.inspect.get("Config")
        if cfg is None:
            return default
        return cfg.get(key, default)

    def _get_cmd(self):
        return self._getconfig("Cmd", [ "/bin/sh" ])
        
    def _interactive(self):
        return (self._getconfig("AttachStdin", False) and
                self._getconfig("AttachStdout", False) and
                self._getconfig("AttachStderr", False))
        
    def _running(self):
        if self._interactive():
            cmd = ["/usr/bin/docker", "exec", "-t", "-i", self.name]
            if self.command:
                cmd += self.command
            else:
                cmd += self._get_cmd()
            return subprocess.check_call(cmd, stderr=subprocess.PIPE)
        else:
            print "Container is running"

    def _start(self):
        if self._interactive():
            if self.command:
                subprocess.check_call(["/usr/bin/docker", "start", self.name], stderr=subprocess.PIPE)
                return subprocess.check_call(["/usr/bin/docker", "exec", "-t", "-i", self.name] + self.command)
            else:
                return subprocess.check_call(["/usr/bin/docker", "start", "-i", "-a", self.name], stderr=subprocess.PIPE)
        else:
            if self.command:
                subprocess.check_call(["/usr/bin/docker", "start", self.name], stderr=subprocess.PIPE)
                return subprocess.check_call(["/usr/bin/docker", "exec", "-t", "-i", self.name] + self.command)
            else:
                return subprocess.check_call(["/usr/bin/docker", "start", self.name], stderr=subprocess.PIPE)

    def _get_args(self, label):
        config = self.inspect["Config"]
        if config and "Labels" in config:
            if config["Labels"] and label in config["Labels"]:
                return config["Labels"][label].split()
        return None

    def run(self):
        try:
            self.inspect = self.d.inspect_container(self.name)
        except docker.errors.APIError:
            pass

        if self.inspect:
            # Container exists
            if self.inspect["State"]["Running"]:
                return self._running()
            else:
                return self._start()

        # Container does not exist
        try:
            self.inspect = self.d.inspect_image(self.image)
        except docker.errors.APIError:
            # Image does not exist
            self.update()
            self.inspect = self.d.inspect_image(self.image)

        if self.spc:
            if self.command:
                args = self.SPC_ARGS + self.command
            else:
                args = self.SPC_ARGS +  self._get_cmd()

            cmd = self.gen_cmd(args)
            print(cmd)
        else:
            missing_RUN = False
            args = self._get_args("RUN")
            if not args:
                missing_RUN = True
                args = self.RUN_ARGS + self._get_cmd()

            cmd = self.gen_cmd(args)
            print(cmd)

            if missing_RUN:
                subprocess.check_call(cmd, env={
                    "CONFDIR": "/etc/%s" % self.name,
                    "LOGDIR": "/var/log/%s" % self.name,
                    "DATADIR":"/var/lib/%s" % self.name}, shell=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
                return self._start()

        subprocess.check_call(cmd, env={
            "CONFDIR": "/etc/%s" % self.name,
            "LOGDIR": "/var/log/%s" % self.name,
            "DATADIR":"/var/lib/%s" % self.name}, shell=True)


    def _rpmostree(self, *args):
        os.execl("/usr/bin/rpm-ostree", "rpm-ostree", *args)

    def host_status(self):
        self._rpmostree("status")

    def host_upgrade(self):
        argv = ["upgrade"]
        if self.args.reboot:
            argv.append("--reboot")
        self._rpmostree(*argv)

    def host_rollback(self):
        argv = ["rollback"]
        if self.args.reboot:
            argv.append("--reboot")
        self._rpmostree(*argv)

    def uninstall(self):
        try:
            self.inspect = self.d.inspect_container(self.name)
            args = self._get_args("UNINSTALL")
            if args:
                cmd = self.gen_cmd(args)
                print(cmd)
                subprocess.check_call(cmd, env={
                    "CONFDIR": "/etc/%s" % self.name,
                    "LOGDIR": "/var/log/%s" % self.name,
                    "DATADIR":"/var/lib/%s" % self.name}, shell=True)
        except docker.errors.APIError:
            pass

        if self.name != self.image:
            try:
                # Attempt to remove container, if it exists just return
                self.d.stop(self.name)
                self.d.remove_container(self.name)
                return
            except:
                # On exception attempt to remove image
                pass

        try:
            self.d.stop(self.image)
            self.d.remove_container(self.image)
        except docker.errors.APIError:
            pass
        try:
            self.inspect = self.d.inspect_image(self.image)
            subprocess.check_call(["/usr/bin/docker", "rmi", self.image])
        except docker.errors.APIError:
            pass


    def gen_cmd(self,cargs):
        args = []
        for c in cargs:
            if c == "IMAGE":
                args.append(self.image)
                continue
            if c == "IMAGE=IMAGE":
                args.append("IMAGE=%s" % self.image)
                continue
            if c == "NAME=NAME":
                args.append("NAME=%s" % self.name)
                continue
            if c == "NAME":
                args.append(self.name)
                continue
            args.append(c)
        return " ".join(args)

    def info(self):
        try:
            labels = self.d.inspect_image(self.image)["Config"]["Labels"]
        except docker.errors.APIError:
            self.update()
            labels = self.d.inspect_image(self.image)["Config"]["Labels"]
        for label in labels:
            print ("%-13s: %s" % (label, labels[label]))

    def install(self):
        try:
            self.inspect = self.d.inspect_image(self.image)
        except docker.errors.APIError:
            self.update()
            self.inspect = self.d.inspect_image(self.image)

        args = self._get_args("INSTALL")
        if args:
            cmd = self.gen_cmd(args)
            print(cmd)

            return(subprocess.check_call(cmd, env={
                "CONFDIR": "/etc/%s" % self.name,
                "LOGDIR": "/var/log/%s" % self.name,
                "DATADIR":"/var/lib/%s" % self.name}, shell=True))

    def help(self):
        if os.path.exists("/usr/bin/rpm-ostree"):
            return _('Atomic Management Tool')
        else:
            return _('Atomic Container Tool')

    def print_spc(self):
        return " ".join(self.SPC_ARGS)

    def print_run(self):
        return " ".join(self.RUN_ARGS)

    def print_install(self):
        return " ".join(self.INSTALL_ARGS) + " /usr/bin/INSTALLCMD"

    def print_uninstall(self):
        return " ".join(self.INSTALL_ARGS) + " /usr/bin/UNINSTALLCMD"

    def defaults(self):
        print(_("""
        Default configuration

        'atomic run' attempts to read the LABEL RUN field in the container image,
        if this field does not exists atom run defaults to the following command:

 %s

        'atomic install' attempts to read the LABEL INSTALL field in the container image,
        if this field does not exists atom install defaults to the following command:

 %s

        These defaults are suggested values for your container images.

        atomic will replace the NAME and IMAGE fields with the name and image specified via the command,  NAME will be replaced with IMAGE if it is not specified.
        """ % (" ".join(self.RUN_ARGS), " ".join(self.INSTALL_ARGS))))

def SetFunc(function):
    class customAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            setattr(namespace, self.dest, function)
    return customAction

class HelpByDefaultArgumentParser(argparse.ArgumentParser):
    def error(self, message):
        self.print_help()
        sys.stderr.write('\nerror: %s\n' % message)
        sys.exit(2)

if __name__ == '__main__':
    atomic=Atomic()
    parser = HelpByDefaultArgumentParser(description=atomic.help())
    subparser = parser.add_subparsers(help=_("Commands"))
    defp = subparser.add_parser("defaults",help=_("list default commands with which Atomic will RUN/INSTALL/UNINSTALL containers"))

    defp.set_defaults(func=atomic.defaults)

    if os.path.exists("/usr/bin/rpm-ostree"):
        hostp = subparser.add_parser("host",help=_("execute Atomic host commands"))
        host_subparser = hostp.add_subparsers(help=_("Host Commands"))
        rollbackp = host_subparser.add_parser("rollback", help=_("revert Atomic to the previously booted tree"))
        rollbackp.set_defaults(func=atomic.host_rollback)
        rollbackp.add_argument("-r", "--reboot", dest="reboot",
                               action="store_true",
                               help=_("Initiate a reboot after rollback is prepared"))

        statusp = host_subparser.add_parser("status", help=_("List information about all deployments"))
        statusp.set_defaults(func=atomic.host_status)
        upgradep = host_subparser.add_parser("upgrade", help=_("List information about all deployments"))
        upgradep.set_defaults(func=atomic.host_upgrade)
        upgradep.add_argument("-r", "--reboot", dest="reboot",
                              action="store_true",
                              help=_("If an upgrade is available, reboot after deployment is complete"))

    infop = subparser.add_parser("info",
                                 help=_("dispaly label information about an image"),
                                 epilog="atomic info attempts to read and display the LABEL information about an image")
    infop.set_defaults(func=atomic.info)
    infop.add_argument("image", help=_("container image"))

    installp = subparser.add_parser("install",
                                    help=_("execute container image install method"),
                                    epilog="atomic install attempts to read the LABEL INSTALL field in the image, if it does not exist atomic will just pull the image on to your machine.  You could add a LABEL INSTALL command to your Dockerfile like: 'LABEL INSTALL %s'" % atomic.print_install() )

    installp.set_defaults(func=atomic.install)
    installp.add_argument("image", help=_("container image"))
    installp.add_argument("command", nargs=argparse.REMAINDER,
                          help=_("execute container image install method."))
    installp.add_argument("-n", "--name", dest="name",
                      default=None,
                      help=_("name of container"))
    runp = subparser.add_parser("run",
                                help=_("execute container image run method."),
                                epilog="atomic run defaults to the following command, if image does not specify LABEL RUN\n'%s'" % atomic.print_run() )
    runp.set_defaults(func=atomic.run)
    runp.add_argument("-n", "--name", dest="name",
                      default=None,
                      help=_("name of container"))
    runp.add_argument("--spc",
                      default=False,
                      action="store_true",
                      help=_("use super privileged container mode: '%s'") % atomic.print_spc() )
    runp.add_argument("image", help=_("container image"))
    runp.add_argument("command", nargs=argparse.REMAINDER,
                      help=_("command to execute within the container"))

    uninstallp = subparser.add_parser("uninstall",
                                      help=_("execute container image uninstall method"),
                                                                          epilog="atomic install attempts to read the LABEL UNINSTALL field in the image, if it does not exist atomic will just pull the image on to your machine.  You could add a LABEL INSTALL command to your Dockerfile like: 'LABEL UNINSTALL %s'" % atomic.print_uninstall() )

    uninstallp.set_defaults(func=atomic.uninstall)
    uninstallp.add_argument("-n", "--name", dest="name",
                            default=None,
                            help=_("name of container"))
    uninstallp.add_argument("image", help=_("container image"))

    updatep = subparser.add_parser("update",help=_("pull latest container image from repository"))
    updatep.set_defaults(func=atomic.update)
    updatep.add_argument("image", help=_("container image"))

    try:
        args = parser.parse_args()
        atomic.set_args(args)
        sys.exit(args.func())
    except ValueError as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
    except IOError as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(0)
    except subprocess.CalledProcessError as e:
        print ""
        sys.exit(e.returncode)
    except docker.errors.DockerException as e:
        sys.stderr.write("%s\n" % str(e))
        sys.exit(1)
