#!/usr/bin/env python3

"""
Loads all puppet files in environment, parse them, and store the
parsed data in the database.
"""

import subprocess
import json
import os
import time

import pyenc
from pyenc.db import db
import pyenc.model as model


def find(path, **kvs):
    """Wrapper around find(1)."""
    cmdline = ['find', path]
    for k, v in kvs.items():
        cmdline.append(f'-{k}')
        cmdline.append(v)
    cmdline.append('-print0')

    cmd = subprocess.run(cmdline, capture_output=True)
    return (f for f in cmd.stdout.split(b'\0') if f)


class PuppetParseError(Exception):
    def __init__(self, code, msg):
        self.code = code
        self.msg = msg

    def __repr__(self):
        return f'PuppetParserError({self.code}, {self.msg})'

    def __str__(self):
        return repr(self)


def puppet_parse(file):
    cmd = subprocess.Popen(
            ['puppet', 'parser', 'dump', '--format', 'json', file],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    if cmd.returncode and cmd.returncode != 0:
        raise PuppetParseError(cmd.returncode, cmd.stderr.read().decode('UTF-8'))
    else:
        json = cmd.stdout.read()

        if (value := cmd.wait()) != 0:
            raise PuppetParseError(value, cmd.stderr.read().decode('UTF-8'))

        return json


def parse_files(files):
    for i, file in enumerate(files):
        try:
            st = os.stat(file)

            last_modify = st.st_mtime
            old_object = model.PuppetFile.query \
                              .where(model.PuppetFile.path == file) \
                              .first()

            if old_object and old_object.last_parse > last_modify:
                # file unchanged since our last parse, skip
                continue

            print(f'{i}/{len(files)}: {file}')

            if old_object:
                m = old_object
            else:
                m = model.PuppetFile(path=file)
            m.last_parse = time.time()
            m.json = puppet_parse(file)

            yield m

        except PuppetParseError as e:
            # TODO cache error
            print('Error:', e)
            continue


def interpret_file(json_data):
    """Find all classes in json-representation of file."""
    top = json_data['^']
    if top[0] == 'class':
        tmp = top[1]['#']
        idx = tmp.index('name')
        return [tmp[idx + 1]]
        # print(tmp[idx + 1])
    elif top[0] == 'block':
        ret_value = []
        for element in top[1:]:
            if element['^'][0] == 'class':
                tmp = element['^'][1]['#']
                idx = tmp.index('name')
                ret_value.append(tmp[idx + 1])
        return ret_value
    else:
        return []




def main():
    app = pyenc.create_app()
    app.app_context().push()

    path = '/var/lib/machines/busting/etc/puppetlabs/code/environments/production'

    files_gen = find(path, type='f', name='*.pp')
    files = [f for f in files_gen]

    try:
        for puppet_file in parse_files(files):
            db.session.add(puppet_file)
    finally:
        db.session.commit()


    try:
        for puppet_file in model.PuppetFile.query.all():
            try:
                class_names = interpret_file(json.loads(puppet_file.json))
                for class_name in class_names:
                    db.session.add(model.PuppetClass(
                        class_name=class_name,
                        comes_from=puppet_file))
            except Exception as e:
                print(e)
                print(f'Failed: {puppet_file.path}')
    finally:
        db.session.commit()

if __name__ == '__main__':
    main()