#!/usr/bin/env python3

# A passthrough tool for docker
import subprocess
import copy
import os
import re
import sys
from optparse import OptionParser, OptionGroup
import tempfile
import jinja2

HERE = os.path.dirname(__file__)
READIES = os.path.abspath(os.path.join(HERE, ".."))
sys.path.insert(0, READIES)
import paella  # noqa: F401

#----------------------------------------------------------------------------------------------

# any environment variable with the same name as a variable in opts, overrides.
# In the case of variables that are action="append", we'll override by splitting the string on space.

def set_env_vars_as_opts(parser, opts):
    appends = [a.dest for a in parser.option_list if a.action=='append']
    for key in vars(opts).keys():
        e = os.getenv(key, None)
        if e is not None:
            if key in appends:  # opts Value objects don't support standard assignment
                setattr(opts, key, e.split())
            else:
                try:  # because zero and one are a thing
                    setattr(opts, key, bool(int(e)))
                except ValueError:
                    setattr(opts, key, e)

#----------------------------------------------------------------------------------------------

class DockerBuilder(object):
    def __init__(self, opts, args):
        self.VERBOSE = opts.VERBOSE
        self.NOP = opts.NOP
        self.KEEP = opts.KEEP

        self.template = opts.TEMPLATE
        self.include_dirs = [] if opts.INCDIRS is None else opts.INCDIRS
        system_includes = os.path.abspath(os.path.join(READIES, "templates", "dockers"))
        self.include_dirs += [os.path.abspath(os.path.dirname(self.template)), os.getcwd(), system_includes]

        self.temp_dockerfile = opts.DOCKERFILE is None
        if self.temp_dockerfile:
            self.dockerfile = tempfile.mkstemp(prefix='Dockerfile.')[1]
        else:
            self.dockerfile = opts.DOCKERFILE
        self.docker_tag = opts.DOCKER_TAG #.lower()
        self.docker_extra_tags = opts.DOCKER_EXTRA_TAGS #[s.lower() for s in opts.DOCKER_EXTRA_TAGS]
        self.build_dir = opts.BUILD_DIR
        self.docker_options = [] if opts.DOCKER_OPTS is None else opts.DOCKER_OPTS.split()

        self.set_template_variables(opts, args)

    def runner(self, cmdline):
        """A run wrapper, so that we can optionally run commands, but do so identically."""
        if self.NOP:
            sys.stdout.write(' '.join(cmdline) + "\n")
            return None

        if self.VERBOSE:
            sys.stdout.write(' '.join(cmdline) + "\n")
            return subprocess.call(cmdline)

        proc = subprocess.Popen(cmdline, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        while True:
            output = proc.stdout.readline()
            if proc.poll() is not None:
                break
            if output:
                sys.stdout.write(output.strip().decode('utf-8') + "\n")
        rc = proc.poll()
        return rc

    def add_env_var(self, var):
        n = str.strip(var)
        val = os.environ.get(var)
        if val is not None:
            self.variables[n] = val

    def add_var_exp(self, exp):
        m = re.match(r'^([^=\s]+)=(.*)', exp)
        if not m is None:
            n = m.group(1)
            v = m.group(2)
            self.variables[n] = v
        else:
            self.add_env_var(exp)

    def set_template_variables(self, opts, args):
        self.variables = {}

        if opts.USE_ALL_ENV_VARS:
            self.variables.update(copy.copy(os.environ))

        if opts.ENV_VAR_PREFIX:
            self.variables.update({k: v for k, v in os.environ.items() if not k.find(opts.ENV_VAR_PREFIX)})

        if not opts.ENV_VARS is None:
            for n in opts.ENV_VARS:
                self.add_env_var(n)

        if not opts.VARS is None:
            for d in opts.VARS:
                self.add_var_exp(d)

        # now, you can pass X=Y on the command line, like a regular argument
        for a in args:
            try:
                k, v = a.split("=")
                if v != "":
                    self.variables[k] = v
            except ValueError:
                pass

    def generate(self):
        """Generate the docker file
        template     - The source file to use, when generating the destination
        dockerfile   - Output destination of the generated context
        include_dirs - A list of directories, containing source templates to optionally import
        variables    - Context: dictionary used when rendeeing the template
        """
        loader = jinja2.FileSystemLoader(self.include_dirs)
        env = jinja2.Environment(loader=loader)
        tmpl = loader.load(name=self.template, environment=env)

        generated = tmpl.render(self.variables)
        with open(self.dockerfile, "w+") as file:
            file.write(generated)

    def build(self):
        """Builds and tags the docker image"""
        ENV['BUILDKIT_PROGRESS'] = 'plain'
        build_cmd = ["docker", "build", "-t", self.docker_tag, "-f", self.dockerfile] + self.docker_options + [self.build_dir]
        res = self.runner(build_cmd)
        if self.temp_dockerfile and not self.KEEP:
            os.unlink(self.dockerfile)
        if res == 0 or self.NOP:
            for tag in self.docker_extra_tags:
                self.runner(["docker", "tag", self.docker_tag, tag])
        return res

    def publish(self):
        """Pushes the docker images upstream."""
        for tag in [self.docker_tag] + self.docker_extra_tags:
            self.runner(["docker", "push", tag])


#----------------------------------------------------------------------------------------------

HELPTEXT = r'''
                                ##         .         
                          ## ## ##        ==         
                       ## ## ## ## ##    ===         
                   /"""""""""""""""""\___/ ===       
              ~~~ {~~ ~~~~ ~~~ ~~~~ ~~~ ~ /  ===- ~~~
                   \______ o           __/           
______           _   \    \         __/              
|  _  \         | |   \____\_______/                 
| | | |___   ___| | _____ _ __ __ _ _ __ ___   __ _  
| | | / _ \ / __| |/ / _ \ '__/ _` | '_ ` _ \ / _` | 
| |/ / (_) | (__|   <  __/ | | (_| | | | | | | (_| | 
|___/ \___/ \___|_|\_\___|_|  \__,_|_| |_| |_|\__,_| 
'''

class OptionParser1(OptionParser):
    def print_help(self):
        print(HELPTEXT)
        super().print_help()

if __name__ == "__main__":
    parser = OptionParser1("Usage: %prog [options] [VAR=value ...]")

    parser.add_option('-n', '--nop', dest='NOP', action='store_true',
                      help="Just print commands, don't execute")
    parser.add_option('-v', '--verbose', dest='VERBOSE', action='store_true',
                      help="Echo commands being run")

    # template variables options
    vars_opts = OptionGroup(parser, 'Template variables')
    parser.add_option_group(vars_opts)
    vars_opts.add_option('-d', dest='VARS', action='append', metavar=('VAR=value'), 
                         help='Define variable for template (repeatable)')
    vars_opts.add_option('-e', dest='ENV_VARS', action='append', metavar=('VAR'),
                         help="Environment variable to be passed to the template (repeatable)")
    vars_opts.add_option('-E', '--env-vars-all', dest='USE_ALL_ENV_VARS', action='store_true',
                         help="If set, all environment variables will be passed to the template")
    vars_opts.add_option('--env-prefix', dest='ENV_VAR_PREFIX', metavar=('PREFIX'), type="str",
                         help="If set, environment variables with this prefix will be used as variables")

    # tempalte generation options
    gen_opts = OptionGroup(parser, 'Template generation')
    parser.add_option_group(gen_opts)
    gen_opts.add_option('-s', '--src', dest='TEMPLATE', type="str", metavar="file", default="dockerfile.tmpl",
                        help="Dockerfile template")
    gen_opts.add_option('-i', '--include', dest='INCDIRS', action="append", metavar="directory",
                        help="Directories of template files to include")
    gen_opts.add_option('-f', '--file', dest='DOCKERFILE', action='store', metavar="file", default=None,
                        help="Dockerfile to generate. If not set, a temp file will be created.")
    gen_opts.add_option('-g', '--generate-and-exit', dest="GENERATE_ONLY", action="store_true",
                        help="Set, to generate the dockerfile, then exit")
    gen_opts.add_option('-k', '--keep', dest='KEEP', action='store_true',
                        help="Keep generated file (when used with a temp file)")

    # docker build options
    build_opts = OptionGroup(parser, 'Build options')
    parser.add_option_group(build_opts)
    build_opts.add_option('--build-dir', dest='BUILD_DIR', action='store', type="str", default='.',
                          help='Directory in which to build')
    build_opts.add_option('-D', '--dockeropts', dest='DOCKER_OPTS', type="str",
                          help="A quoted string, with options to pass to docker build - if building")
    build_opts.add_option('-t', '--tag', dest='DOCKER_TAG', type="str", default=None,
                          help="The named tag to build for the docker image")
    build_opts.add_option('-T', '--extra-tags', dest='DOCKER_EXTRA_TAGS', action="append", default=[],
                          help="An appendable list of docker tags to push")

    # docker publish options
    pub_opts = OptionGroup(parser, 'Publish options')
    parser.add_option_group(pub_opts)
    pub_opts.add_option('-P', '--publish', dest='DOCKER_PUBLISH', action="store_true", default=False,
                        help="Set, if you want to push dockers to dockerhub, after the build")
    pub_opts.add_option('-p', '--publish-only', dest='DOCKER_PUBLISH_ONLY', action="store_true", default=False,
                        help="If set, push the built tags. Nothing is built or generated")
    # pub_opts.add_option('--force-push', dest='DOCKER_FORCE_PUBLISH', action='store_true', default=False,
    #                     help='Push the branch upstream - even if it does not match the version strategy')

    opts, args = parser.parse_args()
    if len(sys.argv) == 1:
        parser.print_help()
        exit(0)

    set_env_vars_as_opts(parser, opts)

    if not os.path.isfile(opts.TEMPLATE):
        sys.stderr.write("Template file '{}' does not exist.\n".format(opts.TEMPLATE))
        sys.exit(3)

    if not opts.DOCKER_TAG:
        sys.stderr.write("Docker tag is missing.\n")
        sys.exit(3)
        
    db = DockerBuilder(opts, args)

    db.generate()
    if opts.GENERATE_ONLY:
        sys.stderr.write("Wrote generated file to {}. Exiting.\n".format(opts.DOCKERFILE))
        sys.exit(0)

    # if we're only pushing, things are already built
    if opts.DOCKER_PUBLISH_ONLY is False:
        r = db.build()
        if opts.NOP is False and r != 0:
            sys.stderr.write("docker build failed, exiting.\n")
            sys.exit(r)

    # publish
    if opts.DOCKER_PUBLISH or opts.DOCKER_PUBLISH_ONLY:
        db.publish()
