#!/usr/bin/env python3

import argparse
import ast
import os

import xlsxwriter


def parse_drivers_dir(options):
    drivers = {}
    for driver in os.listdir(options.cloud_dir):
        if not driver.endswith('.py') or driver == '__init__.py':
            continue
        funcs = parse_driver(os.path.join(options.cloud_dir, driver))
        drivers[driver] = funcs
    return dict(sorted(drivers.items(), key=lambda i: i[0]))


def parse_driver(driver):
    with open(driver, "r") as source:
        tree = ast.parse(source.read())
    functions = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
    func_list = []
    for fn in functions:
        doc = ast.get_docstring(fn)
        doc = doc.splitlines()[0] if doc else ''
        if fn.name.startswith('_'):
            continue
        func_list.append((fn.name, doc))

    assignments = [
        node.targets[0].id for node in tree.body
        if isinstance(node, ast.Assign) and
        isinstance(node.value, ast.Call) and
        isinstance(node.value.func, ast.Name) and
        node.value.func.id == 'namespaced_function'
    ]

    for assignment in assignments:
        func_list.append((assignment, 'NAMESPACED FUNCTION'))
    func_list.sort()
    return func_list


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        '--cloud-dir',
        type=str,
        help='Path to Salt cloud drivers directory',
        required=True,
    )
    parser.add_argument(
        '--out',
        type=str,
        default='salt-cloud-drivers.xlsx',
        help='Name of the resulting xlsx file',
    )
    parser.add_argument(
        '--group',
        action='store_true',
        help='Group functions by name',
    )
    options = parser.parse_args()
    if not os.path.isdir(options.cloud_dir):
        parser.exit('Cloud dir should be a directory: {}'.format(options.cloud_dir))

    drivers = parse_drivers_dir(options)

    workbook = xlsxwriter.Workbook(options.out)
    worksheet = workbook.add_worksheet()

    if options.group:
        all_funcs = sorted(list(set([f[0] for k, v in drivers.items() for f in v])))
        for fi, fn in enumerate(all_funcs):
            worksheet.write(fi + 1, 0, fn)
        for drv_index, (drv_name, functions) in enumerate(drivers.items()):
            worksheet.write(0, drv_index + 1, drv_name)
            for fi, fn in enumerate(functions):
                fn_index = all_funcs.index(fn[0])
                worksheet.write(fn_index + 1, drv_index + 1, fn[0])
                worksheet.write_comment(fn_index + 1, drv_index + 1, fn[1])
    else:
        for drv_index, (drv_name, functions) in enumerate(drivers.items()):
            worksheet.write(0, drv_index + 1, drv_name)
            for fi, fn in enumerate(functions):
                worksheet.write(fi + 1, drv_index + 1, fn[0])
                worksheet.write_comment(fi + 1, drv_index + 1, fn[1])

    workbook.close()