#!/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
from ConfigParser import SafeConfigParser
import getpass

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",
        "file" : "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)

def get_working_root(io_params):
    return os.path.commonprefix(map(os.path.realpath, io_params.values()))

def map_to_dockerpath(io_params, docker_working_root="/zetdata"):
    working_root = get_working_root(io_params)
    if working_root in ['', '/']:
        print("WARNING! Please arrange your data into single directory.")
    docker_paths = map(lambda p:os.path.join(docker_working_root, os.path.relpath(p, working_root)), io_params.values())
    working_volume_str = "%s:%s" % (working_root, docker_working_root)
    return working_volume_str, " ".join(["%s=%s" % (k,os.path.join(docker_working_root, os.path.relpath(v, working_root))) for k,v in io_params.viewitems()])

def print_spec_json(obj):
    print("Name         : %s(%s)" % (obj['Name'], obj['Version']))
    print("CategoryTags : %s" % obj['CategoryTags'])
    print("Params")
    for k,v in obj['Param'].items():
        print("  %s(%s) : '%s'" % (k,v['Type'], v['Default']))
    print("Inputs")
    for k,v in obj['Input'].items():
        print("  %s(%s)" % (k, ",".join(v)))
    print("Outputs")
    for k,v in obj['Output'].items():
        print("  %s(%s)" % (k, ",".join(v)))

def gen_dummy_global_param(username=None):
    if not username:
        username = getpass.getuser()
    return {
        "userName" : {
            "Type" : "string",
            "Val" : username
        },
        "userId" : {
            "Type" : "string",
            "Val" : "123"
        },
        "jobId" : {
            "Type" : "string",
            "Val" : "456"
        },
        "blockId" : {
            "Type" : "string",
            "Val" : "789"
        },
        "hue_server" : {
            "Type" : "string",
            "Val" : "http://192.168.1.20:8888/"
        }
    }

class ZetModule(object):
    def __init__(self, username, module_home=None, keep_files=False):
        self.module_home = os.path.abspath(module_home or '.')
        self.username = username
        self.keep_files = keep_files
        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']

param_types = ['integer', 'enum', 'float', 'string', 'file']

def gen_base_image_option(module_type):
    if module_type == 'basic':
        default_base_image = "zetdata/ubuntu:trusty"
    elif module_type == 'hive':
        default_base_image = "zetdata/cdh:4"
    elif module_type == 'emr_hive':
        default_base_image = "zetdata/ubuntu:trusty"
    else:
        default_base_image = "zetdata/ubuntu:trusty"
    
    return click.Option(('--base-image', '-b'), prompt="Base Image",
                        type=click.Choice(base_images),
                        default=default_base_image,
                        help="Base Image")

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

class InitCLI(click.MultiCommand):
    def list_commands(self, ctx):
        rv = ['basic', 'hive', 'pig', 'emr_hive', 'emr_pig']
        return rv

    def get_command(self, ctx, module_type):

        @click.pass_context
        def init_callback(ctx, name, description, version, cmd, base_image):
            print("init %s" % module_type)
            obj = OrderedDict()
            obj['Name'] = name
            obj['Description'] = description
            obj['Version'] = version
            obj['Cmd'] = cmd
            obj['Param'] = {}
            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/%s' % module_type))

            for tmpl_file in env.list_templates():
                target_file = os.path.splitext(tmpl_file)[0]
                tmpl = env.get_template(tmpl_file)
                with open(os.path.join(target_path, target_file), "w") as f:
                    f.write(tmpl.render(obj))

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

        params = [
            click.Option(('--name','-n'), prompt="Module Name", required=True,
                         help="Module Name"),
            click.Option(('--description', '-d'), prompt="Module Description", required=True,
                         help="Module Description"),
            click.Option(('--version', '-v'), prompt="Module Version",
                         default="0.1",
                         help="Module Version"),
            click.Option(('--cmd', '-c'), prompt="Module Entry Command",
                         default="/usr/bin/python main.py",
                         help="Entry Command"),
            gen_base_image_option(module_type)
        ]

        return click.Command(module_type, help="Create a '%s' type of module" % module_type,
                             params=params, callback=init_callback)

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

@cli.command(short_help="Add a 'Param' to 'spec.json'")
@click.argument('param_keys', nargs=-1)
@click.argument('param_type', nargs=1, required=True)
@click.pass_context
def param_add(ctx, param_keys, param_type):
    data = safe_get_spec_json(ctx)
    for k in param_keys:
        data['Param'][k] = { 'Default' : '', 'Type': param_type }
    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))
    print_spec_json(data)

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

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

@cli.command(short_help="Add a 'Input' parameter to 'spec.json'")
@click.argument('input_name', nargs=1)
@click.argument('input_types', nargs=-1, required=True)
@click.pass_context
def input_add(ctx, input_name, input_types):
    data = safe_get_spec_json(ctx)
    data['Input'][input_name] = list(input_types)
    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))
    print_spec_json(data)

@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=(',', ': ')))
    print_spec_json(data)

@cli.command(short_help="Add a 'Output' parameter to 'spec.json'.")
@click.argument('output_name', nargs=1)
@click.argument('output_types', nargs=-1, required=True)
@click.pass_context
def output_add(ctx, output_name, output_types):
    data = safe_get_spec_json(ctx)
    data['Output'][output_name] = list(output_types)
    with open("spec.json", "w") as sj_out:
        sj_out.write(json.dumps(data, indent=4, separators=(',', ': ')))
    print_spec_json(data)

@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=(',', ': ')))
    print_spec_json(data)

@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 RunCLI(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['Param'].iteritems():
            params.append(click.Option(("--param-%s" % k, ), prompt="Param '%s'"%k, default=v['Default'], type=gettype(v['Type']), help="Param(%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"))

        def build_zetrt(kwargs, zetrt_dir="."):
            param = {re.sub(r'^param_(.*)', r'\1', k):{"Type":"string", "Val":v} for k,v in kwargs.viewitems() if re.match(r'^param_(.*)', k)}

            obj = { "PARAM" : param, "GLOBAL_PARAM" : gen_dummy_global_param()}
            import tempfile
            zetrt_file = tempfile.NamedTemporaryFile(mode="w+", suffix=".json", prefix="./zetrt_tmp_", dir=zetrt_dir, delete=False)
            zetrt_file.write(json.dumps(obj, indent=4, separators=(',', ': ')))
            zetrt_file.close()
            return zetrt_file

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

            # split params into two groups
            io_params = {k:v for k,v in kwargs.viewitems() if not re.match(r'^param_(.*)', k)}
            io_params_str = " ".join(["%s=%s" % (k,v) for k,v in io_params.viewitems()])
            zetrt_file = build_zetrt(kwargs)

            # Build command to execute
            print("Running in local...")
            cmd = "ZETRT=%s %s %s" % (zetrt_file.name, spec_json['Cmd'], io_params_str)
            print("Executing : '%s'" % cmd)
            ret = os.system(cmd)
            if not ctx.obj.keep_files:
                os.remove(zetrt_file.name)
            sys.exit(ret)

        @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
            io_params = {k:v for k,v in kwargs.viewitems() if not re.match(r'^param_(.*)', k)}
            io_vol_str, io_params_str = map_to_dockerpath(io_params)
            zetrt_file = build_zetrt(kwargs)
            zetrt_docker_filename = os.path.join("/home/work/", os.path.relpath(zetrt_file.name))
            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...")
            cur_dir = os.path.realpath(os.path.curdir)
            cmd = "docker run -i -v %s:/home/work/ -v %s -w=/home/run -e ZETRT=%s -t %s %s %s" % (cur_dir, io_vol_str, zetrt_docker_filename, module_path, spec_json['Cmd'], io_params_str)
            print("Executing : '%s'" % cmd)

            ret = os.system(cmd)
            if not ctx.obj.keep_files:
                os.remove(zetrt_file.name)
            sys.exit(ret)

        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
    jinja2_env = Environment(loader=PackageLoader('screwjack', 'misc_templates'))
    template = jinja2_env.get_template("draw_spec_json.dot.j2")
    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())
    def _build():
        build_cmd = 'docker build --no-cache=true -t %s .' % module_path
        print("Executing: %s" % build_cmd)
        if os.system(build_cmd) != 0:
            print("Failed to build '%s'" % module_path)
            sys.exit(-1)

    img_date_str = commands.getoutput('docker inspect -f "{{ .created }}" %s' % module_path)
    print(img_date_str)
    try:
        img_date = parse(img_date_str)
    except Exception as e:
        # No such image
        print("Image '%s' not found, force to rebuild" % module_path)
        _build()
        ctx.exit()

    if force:
        _build()
        ctx.exit()
    modified_files = list(files_in_images(img_date))
    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=RunCLI, 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__":
    if 'DATACANVAS_USERNAME' in os.environ:
        cli(default_map={'username' : os.environ['DATACANVAS_USERNAME']})
    else:
        cfg = SafeConfigParser()
        cfg.read(os.path.expanduser("~/.screwjack.cfg"))
        if cfg.has_section('user') and cfg.has_option('user', 'username'):
            cli(default_map={'username' : cfg.get('user', 'username')})
        else:
            cli()
