#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""A command line tool for manipulating docker modules.
"""

import click
import json
import os
import re
import sys
from collections import OrderedDict
from itertools import izip, chain
from dateutil.parser import parse
from datetime import datetime
import time
import commands
from pytz import UTC

try:
    import screwjack
except:
    sys.path.append(os.path.join(os.path.dirname(__file__), "../"))
    import screwjack

def gettype(name):
    type_map = {
        "string" : "str",
        "integer" : "int",
        "float" : "float",
        "enum" : "str"
    }
    if name not in type_map:
        raise ValueError(name)
    name = type_map[name]
    t = getattr(__builtins__, name)
    if isinstance(t, type):
        return t
    raise ValueError(name)

def safe_get_spec_json(ctx):
    if not ctx.obj.spec_json:
        print("Could not find 'spec.json' in current directory.")
        ctx.exit()
    return ctx.obj.spec_json

def check_docker_image(image_name):
    from subprocess import Popen, PIPE
    p = Popen('docker inspect -f "{{ .id }}" %s' % image_name,
              shell=True, stdout=PIPE, stderr=PIPE, close_fds=True)
    return (p.wait() == 0)

class ZetModule(object):
    def __init__(self, username, module_home=None):
        self.module_home = os.path.abspath(module_home or '.')
        self.username = username
        sj_filename = os.path.join(module_home, "spec.json")
        if not os.path.isfile(sj_filename):
            self.spec_json = None
        else:
            with open(sj_filename, "r") as sj_in:
                self.spec_json = json.load(sj_in, object_pairs_hook=OrderedDict)

base_images = ['zetdata/ubuntu:trusty', 'zetdata/ubuntu:14.04',
               'zetdata/ubuntu:saucy', 'zetdata/ubuntu:13.10',
               'zetdata/ubuntu:raring', 'zetdata/ubuntu:13.04',
               'zetdata/ubuntu:precise', 'zetdata/ubuntu:12.04',
               'zetdata/ubuntu:lucid', 'zetdata/ubuntu:10.4',
               'zetdata/sci-python:2.7', 'zetdata/cdh:4']

@click.group()
@click.option('--username', envvar='DATACANVAS_USERNAME', required=True)
@click.option('--module-home', envvar="DATACANVAS_MODULE_HOME", default=".")
@click.pass_context
def cli(ctx, username, module_home):
    ctx.obj = ZetModule(username, module_home)

@cli.command(short_help="Create a module.")
@click.option('--name', prompt="Module Name", required=True)
@click.option('--description', prompt="Module Description", required=True)
@click.option('--version', prompt="Module Version",
              default="0.1")
@click.option('--cmd', prompt="Module Entry Command",
              default="/usr/bin/python main.py")
@click.option('--base-image', prompt="Base Image",
              type=click.Choice(base_images),
              default="zetdata/ubuntu:trusty")
def init(name, description, version, cmd, base_image):
    obj = OrderedDict()
    obj['Name'] = name
    obj['Description'] = description
    obj['Version'] = version
    obj['Cmd'] = cmd
    obj['Env'] = {}
    obj['Input'] = {}
    obj['Output'] = {}
    obj['BaseImage'] = base_image

    target_path = obj['Name'].lower()
    if os.path.exists(target_path):
        print("Path %s exist, can not create" % target_path)
        exit(-1)

    # Generate files
    os.makedirs(target_path)

    from jinja2 import Environment, PackageLoader
    env = Environment(loader=PackageLoader('screwjack', 'templates'))

    src_templates = ["Dockerfile", "spec.json", "specparser.py", "main.py"]

    for tmpl_file in src_templates:
        tmpl = env.get_template(tmpl_file)
        with open(os.path.join(target_path, tmpl_file), "w") as f:
            f.write(tmpl.render(obj))

    # Show Info
    print("Sucessfully created '%s'" % target_path)

@cli.command(short_help="Add a series 'Env' parameters to 'spec.json'")
@click.argument('env_keys', nargs=-1)
@click.pass_context
def env_add(ctx, env_keys):
    data = safe_get_spec_json(ctx)
    for k in env_keys:
        data['Env'][k] = { 'default' : '', 'type': 'string' }
    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))

@cli.command(short_help="Remove a 'Env' parameter from 'spec.json'")
@click.argument('env_key', nargs=1)
@click.pass_context
def env_del(ctx, env_key):

    data = safe_get_spec_json(ctx)
    data['Env'].pop(env_key, 0)

    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))

@cli.command(short_help="Add a series 'Input' parameters to 'spec.json'")
@click.argument('input_args', nargs=-1)
@click.pass_context
def input_add(ctx, input_args):
    input_keys = input_args[::2]
    input_vals = input_args[1::2]

    data = safe_get_spec_json(ctx)
    for k,v in izip(input_keys, input_vals):
        data['Input'][k] = v
    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))

@cli.command(short_help="Remove a 'Input' parameter from 'spec.json'")
@click.argument('input_key', nargs=1)
@click.pass_context
def input_del(ctx, input_key):
    data = safe_get_spec_json(ctx)
    data['Input'].pop(input_key, 0)

    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))

@cli.command(short_help="Add a series 'Output' parameters to 'spec.json'.")
@click.argument('output_args', nargs=-1)
@click.pass_context
def output_add(ctx, output_args):
    input_keys = output_args[::2]
    input_vals = output_args[1::2]

    data = safe_get_spec_json(ctx)
    for k,v in izip(input_keys, input_vals):
        data['Output'][k] = v
    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))

@cli.command(short_help="Remove a 'Output' parameter from 'spec.json'.")
@click.argument('output_key', nargs=1)
@click.pass_context
def output_del(ctx, output_key):
    data = safe_get_spec_json(ctx)
    data['Output'].pop(output_key, 0)

    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))

@cli.command(short_help="Package current module into a tar file.")
@click.pass_context
def package(ctx):
    internal_package()

@cli.command(short_help="Submit current module to spec_server.")
@click.option('--creator-id', prompt="Spec creator id", required=True,
              default=1)
@click.option('--spec-server', prompt="Spec Server URL", required=True,
              default="http://127.0.0.1:3000/spec/push?creator=1")
@click.pass_context
def submit(ctx, creator_id, spec_server):
    import requests
    sj = safe_get_spec_json(ctx)
    filename = "%s-%s.tar" % (sj['Name'].lower(), sj['Version'])
    if not os.path.exists(filename):
        internal_package()

    import urlparse
    r = requests.post(spec_server,
                      files={'moduletar': open(filename, "rb")})
    if r.status_code != 200:
        print("ERROR : Failed to submit")
        print(r.text)
        print(spec_server)
    else:
        print("Sucessful submit module %s" % filename)

class MyCLI(click.MultiCommand):
    def list_commands(self, ctx):
        rv = ['local', 'docker']
        rv.sort()
        return rv

    def get_command(self, ctx, name):
        spec_json = safe_get_spec_json(ctx)

        ns = {}
        params = []

        for k,v in spec_json['Env'].iteritems():
            params.append(click.Option(("--env-%s" % k, ), prompt="Env '%s'"%k, default=v['default'], type=gettype(v['type']), help="Env(%s)" % v['type']))
        for k,v in spec_json['Input'].iteritems():
            params.append(click.Option(("--%s" % k, ), prompt="Input '%s'"%k, help="Input"))
        for k,v in spec_json['Output'].iteritems():
            params.append(click.Option(("--%s" % k, ), prompt="Output '%s'"%k, help="Output"))

        @click.pass_context
        def run_callback(ctx, *args, **kwargs):
            spec_json = safe_get_spec_json(ctx)

            # split params into two groups
            env_params = {re.sub(r'^env_(.*)', r'ZETENV_\1', k):v for k,v in kwargs.viewitems() if re.match(r'^env_(.*)', k)}
            io_params = {k:v for k,v in kwargs.viewitems() if not re.match(r'^env_(.*)', k)}
            env_params_str = " ".join(["%s=%s" % (k,v) for k,v in env_params.viewitems()])
            io_params_str = " ".join(["%s=%s" % (k,v) for k,v in io_params.viewitems()])

            # Build command to execute
            print("Running in local...")
            cmd = "%s %s %s" % (env_params_str, spec_json['Cmd'], io_params_str)
            print("Executing : '%s'" % cmd)
            os.system(cmd)

        @click.pass_context
        def docker_callback(ctx, *args, **kwargs):
            internal_build(ctx, False)
            spec_json = safe_get_spec_json(ctx)

            # split params into two groups
            env_params = {re.sub(r'^env_(.*)', r'ZETENV_\1', k):v for k,v in kwargs.viewitems() if re.match(r'^env_(.*)', k)}
            io_params = {k:v for k,v in kwargs.viewitems() if not re.match(r'^env_(.*)', k)}
            env_params_str = " ".join(["-e %s=%s" % (k,v) for k,v in env_params.viewitems()])
            io_params_str = " ".join(["%s=%s" % (k,v) for k,v in io_params.viewitems()])
            module_path = "%s/%s" % (ctx.obj.username, spec_json['Name'].lower())

            if not check_docker_image(module_path):
                print("ERROR : Can not find image, ")
                print("        please use 'docker build -t %s .'" % module_path)
                print("        to build your image first.")
                ctx.exit()
            else:
                print("Module '%s' found" % module_path)

            # Build command to execute
            print("Running in docker...")
            cmd = "docker run -i -v $(pwd):/home/work -w=/home/run %s -t %s %s %s" % (env_params_str, module_path, spec_json['Cmd'], io_params_str)
            print("Executing : '%s'" % cmd)
            os.system(cmd)

        if name == "local":
            return click.Command(name, params=params, callback=run_callback)
        elif name == "docker":
            return click.Command(name, params=params, callback=docker_callback)
        else:
            return None

@cli.command(short_help="Render current 'spec.json' to a graphviz file")
def draw():
    with open("spec.json", "r") as jf:
        spec_json = json.load(jf)
    from jinja2 import Environment, PackageLoader
    env = Environment(loader=PackageLoader('screwjack', 'templates'))
    template = env.get_template("draw_spec_json.dot")
    print(template.render(spec_json))

@cli.command(short_help="Build current image")
@click.option('--force', is_flag=True, default=False, help='force to rebuild')
@click.pass_context
def build(ctx, force):
    internal_build(ctx, force)

def internal_build(ctx, force):
    spec_json = safe_get_spec_json(ctx)
    module_path = "%s/%s" % (ctx.obj.username, spec_json['Name'].lower())
    img_date = parse(commands.getoutput('docker inspect -f "{{ .created }}" %s' % module_path))
    modified_files = list(files_in_images(img_date))
    def _build():
        build_cmd = 'docker build --no-cache=true -t %s .' % module_path
        print("Executing: %s" % build_cmd)
        os.system(build_cmd)

    if force:
        _build()
        ctx.exit()
    if len(modified_files) > 0:
        print("The following files are modified(against image: '%s'):" % module_path)
        for fn in modified_files:
            print(fn)
        if query_yes_no("Rebuild?"):
            print("Building")
            _build()
    else:
        print("No need for rebuilding.")

@cli.command(cls=MyCLI, short_help="Run module in local/docker mode")
@click.pass_context
def run(ctx, *args, **kvargs):
    pass

def internal_package():
    import re
    files = [i[0][0] for i in [re.findall(r'^ADD (.*) (.*)$', line)
                               for line in open("Dockerfile")]
             if len(i) > 0]
    files.append("Dockerfile")

    with open("spec.json", "r") as sj:
        sj = json.load(sj, object_pairs_hook=OrderedDict)
    filename = "%s-%s.tar" % (sj['Name'].lower(), sj['Version'])

    print("Packaging files: %s into '%s'" % (files, filename))
    import tarfile
    with tarfile.open(filename, "w") as tar:
        for name in files:
            tar.add(name)

def files_in_images(img_date):
    for root, dirs, files in os.walk("."):
        for fn in files:
            p = os.path.join(root, fn)
            file_mtime = UTC.localize(datetime.fromtimestamp(os.path.getmtime(p)))
            if img_date < file_mtime:
                # print("%s : %s" % (img_date, file_mtime))
                yield p

def query_yes_no(question, default="yes"):
    """Ask a yes/no question via raw_input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
        It must be "yes" (the default), "no" or None (meaning
        an answer is required of the user).

    The "answer" return value is one of "yes" or "no".
    """
    valid = {"yes":True,   "y":True,  "ye":True,
             "no":False,     "n":False}
    if default == None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        sys.stdout.write(question + prompt)
        choice = raw_input().lower()
        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            sys.stdout.write("Please respond with 'yes' or 'no' "\
                             "(or 'y' or 'n').\n")

if __name__ == "__main__":
    cli()
