#!/usr/bin/env python3

# Adapted from: https://enix.io/en/blog/pxe-talos/

import argparse
import functools
import json
import pathlib

import requests
import yaml


@functools.cache
def get_schematic_id(schematic: str):
    """Lookup the schematic id associated with a given schematic"""
    r = requests.post("https://factory.talos.dev/schematics", data=schematic)
    r.raise_for_status()
    data = r.json()
    return data["id"]


def schematic_constructor(directory: pathlib.Path):
    """Load specified schematic file and get the assocatied schematic id"""

    def constructor(loader: yaml.SafeLoader, node: yaml.nodes.ScalarNode):
        filename = str(loader.construct_scalar(node))
        try:
            schematic = directory.joinpath(filename).read_text()
            return get_schematic_id(schematic)
        except Exception:
            raise yaml.MarkedYAMLError("Failed to load schematic", node.start_mark)

    return constructor


def get_loader(directory: pathlib.Path):
    """Add special constructors to yaml loader"""
    loader = yaml.SafeLoader
    loader.add_constructor("!schematic", schematic_constructor(directory))

    return loader


@functools.cache
def get_defaults(directory: pathlib.Path, root: pathlib.Path):
    """Compute the defaults from the provided directory and parents."""
    try:
        with open(directory.joinpath("_defaults.yaml")) as fyaml:
            yml_data = yaml.load(fyaml, Loader=get_loader(directory))
    except OSError:
        yml_data = {}

    # Stop recursion when reaching root directory
    if directory != root:
        return get_defaults(directory.parent, root) | yml_data
    else:
        return yml_data


def walk_files(root: pathlib.Path):
    """Get all files that do not start with and underscore"""
    for dirpath, _dirnames, filenames in root.walk():
        for fn in filenames:
            if not fn.startswith("_"):
                yield dirpath.joinpath(fn)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("directory", type=pathlib.Path)
    parser.add_argument("-f", "--filter")
    args = parser.parse_args()

    data = []
    for fullname in walk_files(args.directory):
        filename = (
            str(fullname.relative_to(args.directory).parent) + "/" + fullname.stem
        )

        if args.filter is not None and not filename.startswith(args.filter):
            continue

        with open(fullname) as fyaml:
            yml_data = yaml.load(fyaml, Loader=get_loader(fullname.parent))
        yml_data = get_defaults(fullname.parent, args.directory) | yml_data
        yml_data["hostname"] = fullname.stem
        yml_data["filename"] = filename
        data.append(yml_data)

    # Dump everything to json
    print(json.dumps(data))


if __name__ == "__main__":
    main()