#!/usr/bin/python3

# For description and usage, see the argparse options at the end of the file

import argparse

import os

from sys import exit

import yaml

def read_file_skip_jinja(filename):

    ''' reads a file and returns its content, except lines starting with '{%' '''

    non_jinja_lines = []

    with open(filename) as f:

        for line in f.read().split('\n'):

            if not ('{%' in line or '{{' in line or '{#' in line):

                non_jinja_lines.append(line)

    return '\n'.join(non_jinja_lines)

def get_roles_of_one_minion(minion):

    if not minion.endswith('.sls'):

      if '.' not in minion and '_' not in minion:

        minion = f'{minion}.infra.opensuse.org'

      minion = minion.replace('.', '_') + '.sls'

    file = f'pillar/id/{minion}'

    try:

        content = read_file_skip_jinja(file)

    except FileNotFoundError:

        print(f'File {file} not found.')

        return []

    try:

        roles = yaml.safe_load(content)['roles']

    except KeyError:

        roles = []

    return roles

def get_minions_with_role(role):

    minions = []

    for sls in os.listdir('pillar/id'):

      if sls.endswith('.sls'):

          for item in get_roles_of_one_minion(sls):

              if item == role:

                  minions.append(os.path.splitext(sls)[0].replace('_', '.'))

    return sorted(minions)

def get_roles(with_base=False):

    roles = []

    for sls in os.listdir('pillar/id'):

        if sls == 'README.md':

            continue

        _roles = get_roles_of_one_minion(sls)

        for item in _roles:

            roles.append(item)

    roles = sorted(set(roles))

    if with_base:

      roles = ['base'] + roles

    return roles

def get_roles_including(query):

    """

    Return roles including the specified string in their state SLS file.

    """

    roles = []

    for role in get_roles():

        role2 = role.replace('.', '/')

        file = f'salt/role/{role2}.sls'

        if not os.path.isfile(file):

          file = f'salt/role/{role2}/init.sls'

        role_sls = read_file_skip_jinja(file)

        if query in role_sls:

            roles.append(role)

    return roles

def print_roles():

    parser = argparse.ArgumentParser('Collects all the roles that are assigned to a minion, and returns them as a python array, a yaml list or a plain list (parsable by bash)')

    parser.add_argument('-o', '--out', choices=['bash', 'python', 'yaml'], help='Select different output format. Options: bash (default), python, yaml')

    parser.add_argument('-i', '--including', help='Only print roles including the specified string in their state file.')

    parser.add_argument('-m', '--minion', help='Only print roles assigned to the specified minion.')

    parser.add_argument('-r', '--role', help='Print minions the specified role is assigned to.')

    args = parser.parse_args()

    if args.including and args.minion:

      print('Combining --including and --minion is not supported. But you can send a patch for it. :)')

      exit(1)

    if ( args.including or args.minion ) and args.role:

      print('Combining --role with --including or --minion is not possible.')

      exit(1)

    if args.including:

        roles = get_roles_including(args.including)

    elif args.minion:

        roles = get_roles_of_one_minion(args.minion)

    elif args.role:

        roles = get_minions_with_role(args.role)

    else:

        roles = get_roles()

    if args.out == 'python':

        print(roles)

    elif args.out == 'yaml':

        print(yaml.dump({'roles': roles}, default_flow_style=False))

    else:

        print('\n'.join(roles))

if __name__ == "__main__":

    print_roles()