#!/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

param_type_map = {
    "string" : "str",
    "integer" : "int",
    "float" : "float",
    "enum" : "str",
    "file" : "str"
}

def gettype(name):
    if name not in param_type_map:
        raise ValueError("Invalid type for %s" % name)
    name = param_type_map[name]
    t = getattr(__builtins__, name)
    if isinstance(t, type):
        return t
    raise ValueError(name)

def param_check(param_name, param_type):
    if param_type not in param_type_map.keys():
        raise ValueError("Invalid type for %s" % param_name)
    
    if not re.match("^[A-Za-z0-9_]+$", param_name):
        raise ValueError("Invalid Param name '%s'. Param type must be '[A-Za-z0-9_]*'" % param_name)

def inout_check(inout_name, inout_types, spec_json):
    if not re.match("^[A-Za-z0-9_]+$", inout_name):
        raise ValueError("Invalid Input/Output name '%s'. Input/Output name must '[A-Za-z0-9_]*'" % inout_name)

    if inout_name in spec_json['Input'].keys():
        raise ValueError("Input name '%s' already exist, please choose another one" % inout_name)

    if inout_name in spec_json['Output'].keys():
        raise ValueError("Output name '%s' already exist, please choose another one" % inout_name)

    for t in inout_types:
        if not re.match("^[a-z0-9\._]+$", t):
            raise ValueError("Invalid Input/Output type '%s'. Input/Output type must be '[a-z0-9\._]+$'." % t)

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):
    common_path = os.path.commonprefix(map(os.path.realpath, io_params.values()))
    if os.path.isfile(common_path):
        return os.path.dirname(common_path)
    elif os.path.isdir(common_path):
        return common_path
    else:
        raise Exception("Invalid path type")

def map_to_dockerpath(io_params, docker_working_root="/zetdata"):
    if len(io_params) == 0:
        print("WARNING! None of Inputs or Outputs.")
        return "", ""

    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 = "-v %s:%s" % (working_root, docker_working_root)
    print(" ".join(["%s=%s" % (k,os.path.join(docker_working_root, os.path.relpath(v, working_root))) for k,v in io_params.viewitems()]))
    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, fast_build=True, spec_server=None):
        self.spec_server = spec_server
        self.module_home = os.path.abspath(module_home or '.')
        self.username = username
        self.keep_files = keep_files
        self.fast_build = fast_build
        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 == 'pig':
        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/--no-keep-files', envvar='DATACANVAS_KEEP_FILES', default=False)
@click.option('--fast-build/--full-build', envvar='DATACANVAS_KEEP_FILES', default=True)
@click.option('--spec_server', envvar='DATACANVAS_SPEC_SERVER', required=False)
@click.pass_context
def cli(ctx, username, module_home, keep_files, fast_build, spec_server):
    ctx.obj = ZetModule(username, module_home, keep_files, fast_build, spec_server)

@cli.command(short_help="Show version of Screwjack")
@click.pass_context
def version(ctx):
    import screwjack
    click.echo('Version %s' % screwjack.__version__)
    ctx.exit()

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))

            # TODO:
            if module_type in ['hive', 'emr_hive']:
                os.makedirs(os.path.join(target_path, "resources/files"))
                os.makedirs(os.path.join(target_path, "resources/udfs"))
            if module_type in ['pig', 'emr_pig']:
                os.makedirs(os.path.join(target_path, "resources/udfs"))

            # 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_key', nargs=1)
@click.argument('param_type', nargs=1, required=True)
@click.pass_context
def param_add(ctx, param_key, param_type):
    data = safe_get_spec_json(ctx)
    param_check(param_key, param_type)

    data['Param'][param_key] = { '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)
    inout_check(input_name, input_types, data)

    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)
    inout_check(output_name, output_types, data)

    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="Login to spec_server")
@click.option('--password', prompt=True, hide_input=True)
@click.pass_context
def login(ctx, password):
    if not ctx.obj.username:
        click.echo("Please input username")
        ctx.exit()
    if not ctx.obj.spec_server:
        click.echo("Please input spec_server")
        ctx.exit()

    click.echo("Logining into spec_server:'%s'" % ctx.obj.spec_server)

    cfg = SafeConfigParser()
    cfg.read(os.path.expanduser("~/.screwjack.cfg"))
    if not cfg.has_section('user'):
        cfg.add_section('user')
    cfg.set('user', 'username', ctx.obj.username)
    cfg.set('user', 'spec_server', ctx.obj.spec_server)
    cfg.write(open(os.path.expanduser("~/.screwjack.cfg"), "w"))

    # login into spec_server
    import requests
    spec_server_url = "http://%s/register" % ctx.obj.spec_server
    spec_server_payload = {"user": ctx.obj.username, "passwd": password}
    r = requests.post(spec_server_url, data=json.dumps(spec_server_payload),
                      headers={'Content-type': 'application/json'})
    if r.status_code == 200:
        cfg.set('user', 'spec_auth', r.json()['token'])
        cfg.write(open(os.path.expanduser("~/.screwjack.cfg"), "w"))
    else:
        click.echo("Failed to login into server:")
        print(r.text)

@cli.command(short_help="Submit current module to spec_server.")
@click.pass_context
def submit(ctx):
    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()

    spec_server = "http://%s/spec/push?user=%s" % (ctx.obj.spec_server, ctx.obj.username)

    cfg = SafeConfigParser()
    cfg.read(os.path.expanduser("~/.screwjack.cfg"))

    r = requests.post(spec_server,
                      files={'moduletar': open(filename, "rb")},
                      headers={'x-spec-auth' : cfg.get('user', 'spec_auth')})
    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/ %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="Print summary of 'spec.json'")
@click.pass_context
def show(ctx):
    data = safe_get_spec_json(ctx)
    print_spec_json(data)

@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():
        if ctx.obj.fast_build:
            build_cmd = 'docker build -t %s .' % module_path
        else:
            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()
        return

    if force:
        _build()
        return
    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"))
        default_map = {}
        if cfg.has_section('user') and cfg.has_option('user', 'username'):
            default_map['username'] = cfg.get('user', 'username')
        if cfg.has_section('user') and cfg.has_option('user', 'spec_server'):
            default_map['spec_server'] = cfg.get('user', 'spec_server')

        if default_map:
            cli(default_map=default_map)
        else:
            cli()
