#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cmd
import ansible.runner
from ansible.color import stringc, codeCodes
import ansible.constants as C
from ansible import utils
import ansible.utils.module_docs as module_docs
import sys
import os
import pprint
import pwd
import readline
import rlcompleter
import atexit
from optparse import OptionParser


class colorizer(object):
    def __init__(self, color):
        self.color = color

    def __enter__(self):
        sys.stdout.write("\033[")
        sys.stdout.write(codeCodes[self.color])
        sys.stdout.write("m")

    def __exit__(self, *args):
        sys.stdout.write("\033[0m")


class AnsibleShell(cmd.Cmd):

    ansible = ansible.runner.Runner()
    groups = ansible.inventory.groups_list().keys()
    hosts = ansible.inventory.groups_list()['all']
    modules = []
    serial = 2

    cwd = ''

    def __init__(self, options, args):
        self.options = options
        self.args = args
        self.intro = 'Welcome to the ansible-shell.\nType help or ? to list commands.\n'
        self.set_prompt()
        self.modules = self.list_modules()
        for module in self.modules:
            setattr(self, 'do_' + module, lambda arg, module=module: self.default(module + ' ' + arg))
            setattr(self, 'help_' + module, lambda module=module: self.helpdefault(module))
        cmd.Cmd.__init__(self)

    @staticmethod
    def parse_opts():
        parser = OptionParser()
        parser.add_option("-s", "--sudo", default=False, action="store_true",
                          dest='sudo', help="run operations with sudo (nopasswd)")
        parser.add_option('-U', '--sudo-user', dest='sudo_user',
                          help='desired sudo user (default=root)', default="root")
        parser.add_option('-u', '--user', default=C.DEFAULT_REMOTE_USER,
                          dest='remote_user',
                          help='connect as this user (default=%s)' % C.DEFAULT_REMOTE_USER)

        return parser.parse_args()

    def get_names(self):
        return dir(self)

    def cmdloop(self):
        try:
            cmd.Cmd.cmdloop(self)
        except KeyboardInterrupt:
            self.intro = " "
            self.cmdloop()

    def set_prompt(self):
        self.prompt = stringc(self.options.remote_user + '@/' + self.cwd, 'green')
        if self.cwd in self.groups:
            self.prompt += stringc(' (' + str(len(self.ansible.inventory.groups_list()[self.cwd])) + ')', 'red')
        self.prompt += '[s:' + stringc(str(self.serial), 'green') + ']'
        self.prompt += '$ '

    def list_modules(self):
        modules = []
        for root, dirs, files in os.walk(C.DEFAULT_MODULE_PATH):
            for basename in files:
                modules.append(basename)

        return modules

    def default(self, arg, forceshell=False):
        if not self.cwd:
            print "No host found"
            return False

        if arg.split()[0] in self.modules:
            module = arg.split()[0]
            module_args = ' '.join(arg.split()[1:])
        else:
            module = 'shell'
            module_args = arg

        if forceshell is True:
            module = 'shell'
            module_args = arg

        try:
            opts = self.options
            results = ansible.runner.Runner(
                pattern=self.cwd, forks=self.serial,
                module_name=module, module_args=module_args,
                remote_user=opts.remote_user,
                sudo=opts.sudo, sudo_user=opts.sudo_user
            ).run()
        except:
            return

        if results is None:
            print "No hosts found"
            return False

        for (hostname, result) in results['contacted'].items():
            if 'stderr' in result.keys():
                if not result['stderr']:
                    print "%s\n%s" % (stringc(hostname, 'bright gray'), result['stdout'])
                else:
                    print "%s >>> %s" % (stringc(hostname, 'red'), result['stderr'])
            else:
                if 'failed' not in result.keys():
                    with colorizer('bright gray'):
                        print hostname
                        pprint.pprint(result)
                else:
                    print "%s >>> %s" % (stringc(hostname, 'red'), result)

    def emptyline(self):
        return

    def do_shell(self, arg):
        self.default(arg, True)

    def do_serial(self, arg):
        """Set the number of forks"""
        self.serial = arg
        self.set_prompt()

    def do_cd(self, arg):
        """Change active host/group"""
        if not arg:
            self.cwd = ''
        elif arg == '..':
            try:
                self.cwd = self.ansible.inventory.groups_for_host(self.cwd)[1].name
            except Exception:
                self.cwd = ''
        elif arg == '/':
            self.cwd = ''
        elif arg in self.hosts or arg in self.groups:
            self.cwd = arg
        else:
            print "incorrect path"

        self.set_prompt()

    def do_list(self, arg):
        """List the hosts in the current group"""
        if arg == 'groups':
            items = self.ansible.inventory.list_groups()
        else:
            items = self.ansible.inventory.list_hosts('all' if self.cwd == '' else self.cwd)
        for item in items:
            print item

    def do_sudo(self, arg):
        """Toggle whether plays run with sudo"""
        self.options.sudo = not self.options.sudo
        print "sudo changed to %s" % self.options.sudo

    def do_remote_user(self, arg):
        """Given a username, set the remote user plays are run by"""
        if arg:
            self.options.remote_user = arg
            self.set_prompt()
        else:
            print "Please specify a remote user, e.g. `remote_user root`"

    def do_sudo_user(self, arg):
        """Given a username, set the user that plays are run by when using sudo"""
        if arg:
            self.options.sudo_user = arg
        else:
            print "Please specify a sudo user, e.g. `sudo_user jenkins`"
            print "Current sudo user is %s" % self.options.sudo_user

    def do_EOF(self, args):
        sys.stdout.write('\n')
        return -1

    def do_exit(self, args):
        """Exits from the console"""
        return -1

    def helpdefault(self, module_name):
        if module_name in self.modules:
            in_path = utils.plugins.module_finder.find_plugin(module_name)
            oc, a = ansible.utils.module_docs.get_docstring(in_path)
            print stringc(oc['short_description'], 'bright gray')
            print 'Parameters:'
            for opt in oc['options'].keys():
                print '  ' + stringc(opt, 'white') + ' ' + oc['options'][opt]['description'][0]

    def complete_cd(self, text, line, begidx, endidx):
        mline = line.partition(' ')[2]
        offs = len(mline) - len(text)

        if self.cwd == '':
            completions = self.hosts + self.groups
        else:
            completions = self.ansible.inventory.list_hosts(self.cwd)

        return [s[offs:] for s in completions if s.startswith(mline)]

    def completedefault(self, text, line, begidx, endidx):
        if line.split()[0] in self.modules:
            mline = line.split(' ')[-1]
            offs = len(mline) - len(text)
            completions = self.module_args(line.split()[0])

            return [s[offs:] + '=' for s in completions if s.startswith(mline)]

    def module_args(self, module_name):
        in_path = utils.plugins.module_finder.find_plugin(module_name)
        oc, a = ansible.utils.module_docs.get_docstring(in_path)
        return oc['options'].keys()

    def do_dump(self, a):
        print readline.get_current_history_length()


if __name__ == '__main__':
    # This hack is to work around readline issues on a mac:
    #  http://stackoverflow.com/a/7116997/541202
    if 'libedit' in readline.__doc__:
        readline.parse_and_bind("bind ^I rl_complete")
    else:
        readline.parse_and_bind("tab: complete")
    histfile = os.path.join(os.path.expanduser("~"), ".ansible-shell_history")
    try:
        readline.read_history_file(histfile)
    except IOError:
        pass
    atexit.register(readline.write_history_file, histfile)

    (options, args) = AnsibleShell.parse_opts()
    AnsibleShell(options, args).cmdloop()

