#!/usr/bin/python3

import argparse
import json
import os
import osbuild
import subprocess
import sys


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


def run_stage_interactive(i, name, options, buildroot, input_dir=None, output_dir=None, sit=False):
    print()
    print(f"{RESET}{BOLD}{i}. {name}{RESET} " + json.dumps(options or {}, indent=2))
    print("Inspect with:")
    print(f"\t# nsenter -a --wd=/root -t `machinectl show {buildroot.machine_name} -p Leader --value`")
    print()

    try:
        buildroot.run_stage(name, options, input_dir, output_dir, sit)
    except KeyboardInterrupt:
        print()
        print(f"{RESET}{BOLD}{RED}Aborted{RESET}")
        return False
    except subprocess.CalledProcessError as error:
        print()
        print(f"{RESET}{BOLD}{RED}{name} failed with code {error.returncode}{RESET}")
        return False
    else:
        return True


def run_interactive(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 False

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

    with osbuild.BuildRoot("/run/osbuild") as buildroot:
        for i, stage in enumerate(pipeline["stages"], start=1):
            name = stage["name"]
            options = stage.get("options", {})
            if not run_stage_interactive(i, name, options, buildroot, input_dir=input_dir, sit=sit):
                return False

        assembler = pipeline.get("assembler")
        if assembler:
            name = assembler["name"]
            options = assembler.get("options", {})
            if not run_stage_interactive("A", name, options, buildroot, output_dir=output_dir, sit=sit):
                return False

    return True


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", type=os.path.abspath,
                        help="provide the contents of DIRECTORY to the first stage")
    parser.add_argument("--output", dest="output_dir", metavar="DIRECTORY", type=os.path.abspath,
                        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()

    os.makedirs("/run/osbuild", exist_ok=True)

    if not run_interactive(args.pipeline_path, args.input_dir, args.output_dir, args.sit):
        sys.exit(1)