#!/usr/bin/python3

import argparse
import contextlib
import json
import os
import subprocess
import sys
import tempfile


RESET = "\033[0m"
BOLD = "\033[1m"
RED = "\033[31m"


@contextlib.contextmanager
def tmpfs(save=None):
    """A contextmanager that mounts a tmpfs and returns its location.

    If `save` is given, it should contain the name of an .tar.gz archive to
    which the contents of the tmpfs will be written.
    """

    with tempfile.TemporaryDirectory(prefix="osbuild-tree-", dir=os.getcwd()) as path:
        subprocess.run(["mount", "-t", "tmpfs", "tmpfs", path], check=True)
        try:
            yield path
            if save:
                print(f"Saving tree to {save}...")
                subprocess.run(["tar", "-czf", save, "-C", path, "."], stdout=subprocess.DEVNULL, check=True)
        finally:
            subprocess.run(["umount", path], check=True)

@contextlib.contextmanager
def bindmnt(src_path):
    """A contextmanager that mindmounts a path read-only and returns its location.
    """

    with tempfile.TemporaryDirectory(prefix="osbuild-build-", dir=os.getcwd()) as dst_path:
        subprocess.run(["mount", "-o", "bind,ro", src_path, dst_path], check=True)
        try:
            yield dst_path
        finally:
            subprocess.run(["umount", dst_path], check=True)


def main(pipeline_path, from_archive, save):
    with open(pipeline_path) as f:
        pipeline = json.load(f)

    with tmpfs(save) as tree, bindmnt("/") as root:
        if from_archive:
            r = subprocess.run(["tar", "-xzf", from_archive, "-C", tree])
            if r.returncode != 0:
                return

        for i, stage in enumerate(pipeline["stages"], start=1):
            name = stage["name"]
            options = stage.get("options", {})
            options["tree"] = "/tmp/tree"
            options["state"] = "/tmp/state"

            options_str = json.dumps(options, indent=2)

            r = subprocess.run(["mkdir", "-p", f"{os.getcwd()}/state/{name}"])
            if r.returncode != 0:
                return

            print()
            print(f"{RESET}{BOLD}{i}. {name}{RESET} {options_str}")
            print()

            try:
                subprocess.run(["systemd-nspawn",
                                "--link-journal=no",
                                f"--directory={root}",
                                f"--bind={tree}:/tmp/tree",
                                f"--bind={os.getcwd()}/state/{name}:/tmp/state",
                                f"--bind={os.getcwd()}/stages/{name}:/tmp/stage",
                                "/tmp/stage"],
                    input=options_str, encoding="utf-8", check=True)
            except KeyboardInterrupt:
                print()
                print(f"{RESET}{BOLD}{RED}Aborted{RESET}")
                return 130
            except subprocess.CalledProcessError as error:
                print()
                print(f"{RESET}{BOLD}{RED}{name} failed with code {error.returncode}{RESET}")
                return 1


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build operating system images")
    parser.add_argument("pipeline_path", metavar="PIPELINE",
                        help="json file containing the pipeline that should be built")
    parser.add_argument("--save", metavar="ARCHIVE",
                        help="save the resulting tree to ARCHIVE")
    parser.add_argument("--from", dest="from_archive", metavar="ARCHIVE",
                        help="initialize the tree from ARCHIVE")
    args = parser.parse_args()

    sys.exit(main(**vars(args)))