#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: ken.liu
# @Date:   2017-09-07 10:19:39
# @Last Modified by:   ken.liu
# @Last Modified time: 2017-09-12 16:05:01

import os
import sys
import logging
from collections import namedtuple

from ansible.parsing.dataloader import DataLoader
from ansible.vars import VariableManager

from ansible.inventory import Inventory
from ansible.inventory.group import Group
from ansible.inventory.host import Host

from ansible.playbook.play import Play
from ansible.executor.playbook_executor import PlaybookExecutor
from ansible.executor.task_queue_manager import TaskQueueManager

from ansible.plugins.callback import CallbackBase

logger = logging.basicConfig()

class ResultsCollector(CallbackBase):
    def __init__(self, *args, **kwargs):
        super(ResultsCollector, self).__init__(*args, **kwargs)
        self.host_ok = {}
        self.host_unreachable = {}
        self.host_failed = {}

    def v2_runner_on_unreachable(self, result):
        self.host_unreachable[result._host.get_name()] = result

    def v2_runner_on_ok(self, result, *args, **kwargs):
        self.host_ok[result._host.get_name()] = result

    def v2_runner_on_failed(self, result, *args, **kwargs):
        self.host_failed[result._host.get_name()] = result

class MyInventory(Inventory):
    def __init__(self, resource, loader, variable_manager):
        self.resource = resource
        self.inventory = Inventory(loader=loader, variable_manager=variable_manager, host_list=[])
        self.gen_inventory()

    def gen_inventory(self):
        if isinstance(self.resource, list):
            self.my_add_group(self.resource, 'default_group')
        elif isinstance(self.resource, dict):
            for groupname, hosts_and_vars in self.resource.iteritems():
                self.my_add_group(hosts_and_vars.get('hosts'), groupname, hosts_and_vars.get('vars'))

    def my_add_group(self, hosts, groupname, groupvars=None):
        my_group = Group(name=groupname)

        if groupvars:
            for key, value in groupvars.iteritems():
                my_group.set_variable(key, value)

        for host in hosts:
            hostname = host.get('hostname')
            host_ip = host.get('ip', hostname)
            host_port = host.get('port')
            username = host.get('username')
            password = host.get('password')
            ssh_key = host.get('ssh_key')

            my_host = Host(name=hostname, port=host_port)
            my_host.set_variable('ansible_ssh_host', host_ip)
            my_host.set_variable('ansible_ssh_port', host_port)
            my_host.set_variable('ansible_ssh_user', username)
            my_host.set_variable('ansible_ssh_pass', password)
            my_host.set_variable('ansible_ssh_private_key_file', ssh_key)

            for key, value in host.iteritems():
                if key not in ['hostname', 'port', 'username', 'password']:
                    my_host.set_variable(key, value)

            my_group.add_host(my_host)

        self.inventory.add_group(my_group)

class AnsibleAPI(object):
    def __init__(self, resource, *args, **kwargs):
        self.resource = resource
        self.inventory = None
        self.variable_manager = None
        self.loader = None
        self.options = None
        self.passwords = None
        self.callback = None
        self.__initializeData()
        self.results_raw = {}

    def __initializeData(self):
        Options = namedtuple('Options', [
            'connection', 'module_path', 'forks', 'timeout', 'remote_user',
            'ask_pass', 'private_key_file', 'ssh_common_args', 'ssh_extra_args',
            'sftp_extra_args', 'scp_extra_args', 'become', 'become_method', 'become_user', 'ask_value_pass',
            'verbosity', 'check', 'listhosts', 'listtasks', 'listtags', 'syntax'
            ]
        )

        self.variable_manager = VariableManager()
        self.loader = DataLoader()
        self.options = Options(
            connection='smart',
            module_path='/usr/share/ansible',
            forks=100, timeout=10,
            remote_user='root', ask_pass=False,
            private_key_file='ssh/ansible.pem',
            ssh_common_args=None, ssh_extra_args=None, sftp_extra_args=None, scp_extra_args=None,
            become=None, become_method=None, become_user='root',
            ask_value_pass=False, verbosity=None, check=False, listhosts=False,
            listtasks=False, listtags=False, syntax=False
        )

        self.passwords = dict(sshpass=None, becomepass=None)
        self.inventory = MyInventory(self.resource, self.loader, self.variable_manager).inventory
        self.variable_manager.set_inventory(self.inventory)

    def run(self, host_list, module_name, module_args):
        play_source = dict(
            name='Ansible Play',
            hosts=host_list,
            gather_facts='no',
            tasks=[
                dict(action=dict(module=module_name, args=module_args))
            ]
        )
        play = Play().load(play_source, variable_manager=self.variable_manager, loader=self.loader)

        tqm = None
        self.callback = ResultsCollector()
        try:
            tqm = TaskQueueManager(
                inventory=self.inventory,
                variable_manager=self.variable_manager,
                loader=self.loader,
                options=self.options,
                passwords=self.passwords,
            )
            tqm._stdout_callback = self.callback
            tqm.run(play)
        finally:
            if tqm is not None:
                tqm.cleanup()

    def run_playbook(self, host_list, role_name, role_uuid, temp_param):
        try:
            self.callback = ResultsCollector()
            filenames = ['' + '/handlers/ansible/v1_0/sudoers.yml']
            template_file = ''
            if not os.path.exists(template_file):
                sys.exit()

            host_list_str = ','.join([item for item in host_list])

            extra_vars = {}
            extra_vars['host_list'] = host_list_str
            extra_vars['username'] = role_name
            extra_vars['template_dir'] = template_file
            extra_vars['command_list'] = temp_param.get('cmdList')
            extra_vars['role_uuid'] = 'role-%s' % role_uuid

            self.variable_manager.extra_vars = extra_vars

            executor = PlaybookExecutor(
                playbooks=filenames,
                inventory=self.inventory,
                variable_manager=self.variable_manager,
                loader=self.loader,
                options=self.options,
                passwords=self.passwords,
            )
            executor._tqm._stdout_callback = self.callback
            executor.run()
        except Exception as e:
            print "error:",e.message

    def get_result(self):
        self.results_raw = {'success': {}, 'failed': {}, 'unreachable': {}}

        for host, result in self.callback.host_ok.items():
            self.results_raw['success'][host] = result._result

        for host, result in self.callback.host_failed.items():
            self.results_raw['failed'][host] = result._result.get('msg') or result._result

        for host, result in self.callback.host_unreachable.items():
            self.results_raw['unreachable'][host] = result._result['msg']

        return self.results_raw