#!/usr/bin/env python2

# A tool to check the order of precedence for ansible variables
#
#    TOTAL: 20
#    1. extra_vars
#    2. set_fact
#    3. include_vars
#    4. role_var
#    5. vars_file
#    6. play_var
#    7. pb_host_vars_file
#    8. ini_group_vars_file_parent
#    9. ini_group_vars_file_child
#    10. ini_group_vars_file_all
#    11. ini_host_vars_file
#    12. ini_host
#    13. pb_group_vars_file_child
#    14. pb_group_vars_file_parent
#    15. ini_child
#    16. ini_parent
#    17. ini_all
#    18. pb_group_vars_file_all
#    19. role_parent_default
#    20. role_default

# https://github.com/ansible/ansible/blob/devel/test/integration/test_var_precedence.yml

# bcoca [3:25 PM]
# that looks wrong, ini < pb

# bcoca's version ...
#    1. extra_vars
#    2. set_fact
#    3. include_vars
#    4. role_vars
#    5. vars_file
#    6. play_var
#    7. pb_host_vars_file
#    13. pb_group_vars_file_child
#    14. pb_group_vars_file_parent
#    11. ini_host_vars_file
#    9. ini_group_vars_file_child
#    8. ini_group_vars_file_parent
#    12. ini_host
#    15. ini_child
#    16. ini_parent
#    18. pb_group_vars_file_all
#    10. ini_group_vars_file_all
#    17. ini_all
#    19. role_parent_default
#    20. role_default

import json
import os
import sys
import shutil
import stat
import subprocess
import tempfile
from pprint import pprint
from optparse import OptionParser
from jinja2 import Environment

ENV = Environment()
TESTDIR = tempfile.mkdtemp()


def run_command(args, cwd=None):
    p = subprocess.Popen(
            args,
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE,
            shell=True,
            cwd=cwd
        )
    (so, se) = p.communicate()
    return (p.returncode, so, se)


def clean_test_dir():
    if os.path.isdir(TESTDIR):
        shutil.rmtree(TESTDIR)
    os.makedirs(TESTDIR)


class Role(object):
    def __init__(self, name):
        self.name = name
        self.load = True
        self.dependencies = []
        self.defaults = False
        self.vars = False
        self.tasks = []

    def write_role(self):

        fpath = os.path.join(TESTDIR, 'roles', self.name)
        if not os.path.isdir(fpath):
            os.makedirs(fpath)

        if self.defaults:
            # roles/x/defaults/main.yml
            fpath = os.path.join(TESTDIR, 'roles', self.name, 'defaults')
            if not os.path.isdir(fpath):
                os.makedirs(fpath)
            fname = os.path.join(fpath, 'main.yml')
            with open(fname, 'wb') as f:
                f.write('findme: %s\n' % self.name)

        if self.vars:
            # roles/x/vars/main.yml
            fpath = os.path.join(TESTDIR, 'roles', self.name, 'vars')
            if not os.path.isdir(fpath):
                os.makedirs(fpath)
            fname = os.path.join(fpath, 'main.yml')
            with open(fname, 'wb') as f:
                f.write('findme: %s\n' % self.name)

        if self.dependencies:
            fpath = os.path.join(TESTDIR, 'roles', self.name, 'meta')
            if not os.path.isdir(fpath):
                os.makedirs(fpath)
            fname = os.path.join(fpath, 'main.yml')
            with open(fname, 'wb') as f:
                f.write('dependencies:\n')
                for dep in self.dependencies:
                    f.write('- { role: %s }\n' % dep)
            #import epdb; epdb.st()


class DynamicInventory(object):
    BASESCRIPT = '''#!/usr/bin/python
import json
data = """{{ data }}"""
data = json.loads(data)
print(json.dumps(data, indent=2, sort_keys=True))
'''

    BASEINV = {
        '_meta': {
            'hostvars': {
                'testhost': {}
            }
        }
    }

    def __init__(self, features):
        self.ENV = Environment()
        self.features = features
        self.fpath = None
        self.inventory = self.BASEINV.copy()
        #print('###########################################')
        #pprint(self.inventory)
        #print('###########################################')
        self.build()

    def build(self):
        xhost = 'testhost'
        if 'script_host' in self.features:
            #print([x for x in self.features if 'script' in x])
            self.inventory['_meta']['hostvars'][xhost]['findme'] = 'script_host'
        else:
            self.inventory['_meta']['hostvars'][xhost] = {}

        if 'script_child' in self.features:
            self.inventory['child'] = {
                'hosts': [xhost],
                'vars': {'findme': 'script_child'}
            }

        if 'script_parent' in self.features:

            self.inventory['parent'] = {
                'vars': {'findme': 'script_parent'}
            }

            if 'script_child' in self.features:
                self.inventory['parent']['children'] = ['child']
            else:
                self.inventory['parent']['hosts'] = [xhost]

        if 'script_all' in self.features:
            self.inventory['all'] = {
                'hosts': [xhost],
                'vars': {
                    'findme': 'script_all'
                },
            }
        else:
            self.inventory['all'] = {
                'hosts': [xhost],
            }

        #import epdb; epdb.st()

    def write_script(self):
        fdir = os.path.join(TESTDIR, 'inventory')
        if not os.path.isdir(fdir):
            os.makedirs(fdir)
        fpath = os.path.join(fdir, 'hosts')
        #fpath = os.path.join(TESTDIR, 'inventory')
        self.fpath = fpath

        data = json.dumps(self.inventory)
        t = self.ENV.from_string(self.BASESCRIPT)
        fdata = t.render(data=data)
        with open(fpath, 'wb') as f:
            f.write(fdata + '\n')
        st = os.stat(fpath)
        os.chmod(fpath, st.st_mode | stat.S_IEXEC)
        return fpath

EMPTY_ASSERT_EXPECTED = '''
'''

class VarTestMaker(object):
    def __init__(self, features, dynamic_inventory=False):
        clean_test_dir()
        self.dynamic_inventory = dynamic_inventory
        self.di = None
        self.features = features
        self.inventory = ''
        self.playvars = []
        self.varsfiles = []
        self.playbook = '- hosts: testhost\n'
        self.playbook += '  gather_facts: False\n'
        self.tasks = []
        self.roles = []
        self.ansible_command = None
        self.stdout = None
        self.assert_expected = True

    def write_playbook(self):
        fname = os.path.join(TESTDIR, 'site.yml')
        with open(fname, 'wb') as f:
            f.write(self.playbook)
            if self.playvars:
                f.write('  vars:\n')
                for pv in self.playvars:
                    f.write('    %s\n' % pv)
            if self.varsfiles:
                f.write('  vars_files:\n')
                for vf in self.varsfiles:
                    f.write('    - %s\n' % vf)
            if self.roles:
                f.write('  roles:\n')
                for role in self.roles:
                    #print("writing %s" % role.name)
                    role.write_role()
                    if role.load:
                        f.write('    - %s\n' % role.name)
            if self.tasks:
                f.write('  tasks:\n')
                for task in self.tasks:
                    if isinstance(task, list):
                        for idx,x in enumerate(task):
                            if idx == 0:
                                f.write('    - %s\n' % x)
                            else:
                                f.write('      %s\n' % x)
                    else:
                        f.write('    - %s\n' % task)
                if self.assert_expected:
                    f.write('    - include: assert_expected.yml\n')

    def build(self):

        if self.dynamic_inventory:
            # python based inventory file
            self.di = DynamicInventory(self.features)
            self.inventory_script_path = self.di.write_script()
        else:
            # ini based inventory file
            if 'ini_host' in self.features:
                self.inventory += 'testhost findme=ini_host\n'
            else:
                self.inventory += 'testhost\n'
            self.inventory += '\n'

            if 'ini_child' in self.features:
                self.inventory += '[child]\n'
                self.inventory += 'testhost\n'
                self.inventory += '\n'
                self.inventory += '[child:vars]\n'
                self.inventory += 'findme=ini_child\n'
                self.inventory += '\n'

            if 'ini_parent' in self.features:
                if 'ini_child' in self.features:
                    self.inventory += '[parent:children]\n'
                    self.inventory += 'child\n'
                else:
                    self.inventory += '[parent]\n'
                    self.inventory += 'testhost\n'
                self.inventory += '\n'
                self.inventory += '[parent:vars]\n'
                self.inventory += 'findme=ini_parent\n'
                self.inventory += '\n'

            if 'ini_all' in self.features:
                self.inventory += '[all:vars]\n'
                self.inventory += 'findme=ini_all\n'
                self.inventory += '\n'

            # default to a single file called inventory
            invfile = os.path.join(TESTDIR, 'inventory', 'hosts')
            ipath = os.path.join(TESTDIR, 'inventory')
            if not os.path.isdir(ipath):
                os.makedirs(ipath)

            with open(invfile, 'wb') as f:
                f.write(self.inventory)

        hpath = os.path.join(TESTDIR, 'inventory', 'host_vars')
        if not os.path.isdir(hpath):
            os.makedirs(hpath)
        gpath = os.path.join(TESTDIR, 'inventory', 'group_vars')
        if not os.path.isdir(gpath):
            os.makedirs(gpath)

        if 'ini_host_vars_file' in self.features:
            hfile = os.path.join(hpath, 'testhost')
            with open(hfile, 'wb') as f:
                f.write('findme: ini_host_vars_file\n')

        if 'ini_group_vars_file_all' in self.features:
            hfile = os.path.join(gpath, 'all')
            with open(hfile, 'wb') as f:
                f.write('findme: ini_group_vars_file_all\n')

        if 'ini_group_vars_file_child' in self.features:
            hfile = os.path.join(gpath, 'child')
            with open(hfile, 'wb') as f:
                f.write('findme: ini_group_vars_file_child\n')

        if 'ini_group_vars_file_parent' in self.features:
            hfile = os.path.join(gpath, 'parent')
            with open(hfile, 'wb') as f:
                f.write('findme: ini_group_vars_file_parent\n')

        if 'pb_host_vars_file' in self.features:
            os.makedirs(os.path.join(TESTDIR, 'host_vars'))
            fname = os.path.join(TESTDIR, 'host_vars', 'testhost')
            with open(fname, 'wb') as f:
                f.write('findme: pb_host_vars_file\n')

        if 'pb_group_vars_file_parent' in self.features:
            if not os.path.isdir(os.path.join(TESTDIR, 'group_vars')):
                os.makedirs(os.path.join(TESTDIR, 'group_vars'))
            fname = os.path.join(TESTDIR, 'group_vars', 'parent')
            with open(fname, 'wb') as f:
                f.write('findme: pb_group_vars_file_parent\n')

        if 'pb_group_vars_file_child' in self.features:
            if not os.path.isdir(os.path.join(TESTDIR, 'group_vars')):
                os.makedirs(os.path.join(TESTDIR, 'group_vars'))
            fname = os.path.join(TESTDIR, 'group_vars', 'child')
            with open(fname, 'wb') as f:
                f.write('findme: pb_group_vars_file_child\n')

        if 'pb_group_vars_file_all' in self.features:
            if not os.path.isdir(os.path.join(TESTDIR, 'group_vars')):
                os.makedirs(os.path.join(TESTDIR, 'group_vars'))
            fname = os.path.join(TESTDIR, 'group_vars', 'all')
            with open(fname, 'wb') as f:
                f.write('findme: pb_group_vars_file_all\n')

        if 'play_var' in self.features:
            self.playvars.append('findme: play_var')

        if 'set_fact' in self.features:
            self.tasks.append('set_fact: findme="set_fact"')

        if 'vars_file' in self.features:
            self.varsfiles.append('varsfile.yml')
            fname = os.path.join(TESTDIR, 'varsfile.yml')
            with open(fname, 'wb') as f:
                f.write('findme: vars_file\n')

        if 'include_vars' in self.features:
            self.tasks.append('include_vars: included.yml')
            fname = os.path.join(TESTDIR, 'included.yml')
            with open(fname, 'wb') as f:
                f.write('findme: include_vars\n')

        if 'role_var' in self.features:
            role = Role('role_var')
            role.vars = True
            role.load = True
            self.roles.append(role)

        if 'role_parent_default' in self.features:
            role = Role('role_default')
            role.load = False
            role.defaults = True
            self.roles.append(role)

            role = Role('role_parent_default')
            role.dependencies.append('role_default')
            role.defaults = True
            role.load = True
            self.roles.append(role)

        elif 'role_default' in self.features:
                role = Role('role_default')
                role.defaults = True
                role.load = True
                self.roles.append(role)

        if 'task_vars' in self.features:
            self.tasks.append(
                [
                    'debug: var=findme',
                    'vars:',
                    '   findme: "task_vars"'
                ]
            )
        else:
            self.tasks.append('debug: var=findme')

        if self.assert_expected:
            # write an empty assert_expected for the eval,so copy_local can add a real one
            fname = os.path.join(TESTDIR, 'assert_expected.yml')
            with open(fname, 'wb') as f:
                f.write(EMPTY_ASSERT_EXPECTED)




        self.write_playbook()

    def run(self):
        '''
        if self.dynamic_inventory:
            cmd = 'ansible-playbook -c local -i inventory/hosts site.yml'
        else:
            cmd = 'ansible-playbook -c local -i inventory site.yml'
        '''
        cmd = 'ansible-playbook -c local -i inventory site.yml'
        if 'extra_vars' in self.features:
            cmd += ' --extra-vars="findme=extra_vars"'
        self.ansible_command = cmd
        (rc, so, se) = run_command(cmd, cwd=TESTDIR)
        self.stdout = so

        if rc == 0:
            val = None
            try:
                idx = so.rfind('findme')
                val = so[idx:]
                val = val.split('"')[2]
            except Exception as e:
                print(e)
                import epdb; epdb.st()
            return val
        else:
            print(se)
            import epdb; epdb.st()
            sys.exit(1)
            return False

    def show_tree(self):
        print('## TREE')
        cmd = 'tree %s' % TESTDIR
        (rc, so, se) = run_command(cmd)
        lines = so.split('\n')
        lines = lines[:-3]
        print('\n'.join(lines))
        #import epdb; epdb.st()

    def show_content(self):
        print('## CONTENT')
        cmd = 'find %s -type f | xargs tail -n +1' % TESTDIR
        (rc, so, se) = run_command(cmd)
        print(so)
        #import epdb; epdb.st()

    def show_stdout(self):
        print('## COMMAND')
        print(self.ansible_command)
        print('## STDOUT')
        print(self.stdout)


ASSERT_EXPECTED = '''
- set_fact:
    expected: %s
- assert:
    that:
      - findme == expected
'''


def main():
    #results = []
    features = ['ini_host',
                'ini_child',
                'ini_parent',
                'ini_all',
                'ini_host_vars_file',
                'ini_group_vars_file_parent',
                'ini_group_vars_file_child',
                'ini_group_vars_file_all',
                'pb_group_vars_file_parent',
                'pb_group_vars_file_child',
                'pb_group_vars_file_all',
                'pb_host_vars_file',
                'play_var',
                'role_parent_default',
                'role_default',
                'role_var',
                'include_vars',
                'set_fact',
                'vars_file',
                'task_vars',
                'extra_vars']

    parser = OptionParser()
    parser.add_option('-f', '--feature', action='append')
    parser.add_option('--use_dynamic_inventory', action='store_true')
    parser.add_option('--show_tree', action='store_true')
    parser.add_option('--show_content', action='store_true')
    parser.add_option('--show_stdout', action='store_true')
    parser.add_option('--copy_testcases_to_local_dir', action='store_true')
    (options, args) = parser.parse_args()

    if options.feature:
        for f in options.feature:
            if f not in features:
                print('%s is not a valid feature' % f)
                sys.exit(1)
        features = [x for x in options.feature]

    fdesc = {
        'ini_host': 'host var inside the ini',
        'script_host': 'host var inside the script _meta',
        'ini_child': 'child group var inside the ini',
        'script_child': 'child group var inside the script',
        'ini_parent': 'parent group var inside the ini',
        'script_parent': 'parent group var inside the script',
        'ini_all': 'all group var inside the ini',
        'script_all': 'all group var inside the script',
        'ini_host_vars_file': 'var in inventory/host_vars/host',
        'ini_group_vars_file_parent': 'var in inventory/group_vars/parent',
        'ini_group_vars_file_child': 'var in inventory/group_vars/child',
        'ini_group_vars_file_all': 'var in inventory/group_vars/all',
        'pb_group_vars_file_parent': 'var in playbook/group_vars/parent',
        'pb_group_vars_file_child': 'var in playbook/group_vars/child',
        'pb_group_vars_file_all': 'var in playbook/group_vars/all',
        'pb_host_vars_file': 'var in playbook/host_vars/host',
        'play_var': 'var set in playbook header',
        'role_parent_default': 'var in roles/role_parent/defaults/main.yml',
        'role_default': 'var in roles/role/defaults/main.yml',
        'role_var': 'var in ???',
        'include_vars': 'var in included file',
        'set_fact': 'var made by set_fact',
        'vars_file': 'var in file added by vars_file',
        'task_vars': 'vars defined on the task',
        'extra_vars': 'var passed via the cli'
    }

    dinv = options.use_dynamic_inventory
    if dinv:
        # some features are specific to ini, so swap those
        for idx,x in enumerate(features):
            if x.startswith('ini_') and 'vars_file' not in x:
                features[idx] = x.replace('ini_', 'script_')

    print("TOTAL: %s" % len(features))
    dinv = options.use_dynamic_inventory

    last_removed = None
    index = 1
    while features:
        VTM = VarTestMaker(features, dynamic_inventory=dinv)
        VTM.build()

        if options.show_tree or options.show_content or options.show_stdout:
            print('')

        if options.show_tree:
            VTM.show_tree()

        if options.show_content:
            VTM.show_content()

        res = VTM.run()

        if options.show_stdout:
            VTM.show_stdout()

        if not res or res == 'VARIABLE IS NOT DEFINED!':
            print("ERROR: %s" % res)
            print("tempdir: %s" % TESTDIR)
            print("last_removed: %s" % last_removed)
            print("features: %s" % features)
            sys.exit(1)

        print('%s. %s - %s' % (index, res, fdesc[res]))
        try:
            features.remove(res)
        except Exception as e:
            #import epdb; epdb.st()
            print("ERROR !!!")
            print(e)
            print('features: %s' % features)
            print('res: %s' % res)
            sys.exit(1)
        last_removed = res

        if options.copy_testcases_to_local_dir:

            topdir = 'testcases'

            if index == 1 and os.path.isdir(topdir):
                shutil.rmtree(topdir)

            if not os.path.isdir(topdir):
                os.makedirs(topdir)
            thisindex = str(index)
            if len(thisindex) == 1:
                thisindex = '0' + thisindex
            thisdir = os.path.join(topdir, '%s.%s' % (thisindex, res))
            shutil.copytree(TESTDIR, thisdir)

            set_expected = True
            if set_expected:
                expfile = os.path.join(thisdir, 'assert_expected.yml')
                with open(expfile, 'wb') as f:
                    f.write(ASSERT_EXPECTED % res)

        shutil.rmtree(TESTDIR)
        index += 1

if __name__ == "__main__":
    main()
