#!/usr/bin/python3

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


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


class BuildContainer:
    def __init__(self, path=os.getcwd()):
        self.buildroot = tempfile.mkdtemp(prefix="osbuild-buildroot-", dir=path)
        self.buildroot_mounted = False
        self.tree = tempfile.mkdtemp(prefix="osbuild-tree-", dir=path)
        self.tree_mounted = False
        try:
            subprocess.run(["mount", "-o", "bind,ro", "/", self.buildroot], check=True)
            self.tree_mounted = True
            subprocess.run(["mount", "-t", "tmpfs", "tmpfs", self.tree], check=True)
            self.buildroot_mounted = True
        except subprocess.CalledProcessError:
            self.unmount()
            raise

    def unmount(self):
        if self.tree:
            if self.tree_mounted:
                subprocess.run(["umount", "--lazy", self.tree], check=True)
            os.rmdir(self.tree)
            self.tree = None
        if self.buildroot:
            if self.buildroot_mounted:
                subprocess.run(["umount", "--lazy", self.buildroot], check=True)
            os.rmdir(self.buildroot)
            self.buildroot = None

    def __del__(self):
        self.unmount()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.unmount()


def main(pipeline_path, input_dir, output_dir, sit):
    if output_dir and len(os.listdir(output_dir)) != 0:
        print()
        print(f"{RESET}{BOLD}{RED}Output directory {output_dir} is not empty{RESET}")
        return 1

    with open(pipeline_path) as f:
        pipeline = json.load(f)

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

            # systemd-nspawn silently removes some characters when choosing a
            # machine name from the directory name. The only one relevant for
            # us is '_', because all other characters used by
            # TemporaryDirectory() are allowed. Replace it with 'L's
            # (TemporaryDirectory() only uses lower-case characters)
            machine_name = os.path.basename(container.buildroot).replace("_", "L")

            argv = ["systemd-nspawn",
                    "--as-pid2",
                    "--link-journal=no",
                    "--volatile=yes",
                    f"--machine={machine_name}",
                    f"--directory={container.buildroot}",
                    f"--bind={container.tree}:/tmp/tree",
                    f"--bind={os.getcwd()}/run-stage:/tmp/run-stage",
                    f"--bind={os.getcwd()}/stages/{name}:/tmp/stage",
                    "--bind=/etc/pki"]

            # Optionally pass the input dir to the first stage
            if input_dir:
                options["input_dir"] = "/tmp/input"
                argv.append(f"--bind-ro={os.getcwd()}/{input_dir}:/tmp/input")
            else:
                options["input_dir"] = None

            # Pass the tree r/w to all stages but the last one.
            # The last stage gets the tree ro and optionally gets an output dir
            if i == len(pipeline["stages"]) and output_dir:
                options["output_dir"] = "/tmp/output"
                argv.append(f"--bind={os.getcwd()}/{output_dir}:/tmp/output")

            argv.append("/tmp/run-stage")
            if sit:
                argv.append("--sit")
            argv.append("/tmp/stage")

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

            print()
            print(f"{RESET}{BOLD}{i}. {name}{RESET} {options_str}")
            print("Inspect with:")
            print(f"\t# nsenter -a --wd=/root -t `machinectl show {machine_name} -p Leader --value`")
            print()

            try:
                subprocess.run(argv, 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("--input", dest="input_dir", metavar="DIRECTORY",
                        help="provide the contents of DIRECTORY to the first stage")
    parser.add_argument("--output", dest="output_dir", metavar="DIRECTORY",
                        help="provide the empty DIRECTORY as output argument to the last stage")
    parser.add_argument("--sit", action="store_true",
                        help="keep the build environment up when a stage failed")
    args = parser.parse_args()

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