#!/usr/bin/python3

"""OSTree Image Tests"""

import argparse
import os
import subprocess
import sys
import tempfile
import time

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


class OSBuild:
    def __init__(self, store, outdir):
        self.store = store
        self.outdir = outdir
        self.checkpoints = []

    def run(self, manifest, exports, checkpoints=None):
        cmd = [
            "osbuild",
            "--store", os.fspath(self.store),
            "--output-dir", os.fspath(self.outdir),
            os.fspath(manifest)
        ]

        for checkpoint in self.checkpoints + (checkpoints or []):
            cmd += [
                "--checkpoint", checkpoint
            ]

        for export in exports:
            cmd += [
                "--export", export
            ]

        subprocess.run(cmd, check=True)


def run_tests(args, tmpdir):
    outdir, store = args.output_directory, args.store

    if not outdir:
        outdir = os.path.join(tmpdir, "outdir")
        os.makedirs(outdir)

    if not store:
        store = os.path.join(tmpdir, "store")
        os.makedirs(store)

    print(f"Store at:  {os.path.realpath(store)}")
    print(f"Output at: {os.path.realpath(outdir)}")

    osbuild = OSBuild(store, outdir)

    osbuild.checkpoints = [
        "build",
        "ostree-tree",
        "ostree-commit"
    ]

    print(f"Testing {BOLD}{args.manifest}{RESET}", flush=True)

    path = os.path.join("test", "data", "manifests", args.manifest)

    success = True
    export_names = map(lambda p: p.split("/")[0], args.export)
    osbuild.run(path, export_names)
    for export in args.export:
        path = os.path.join(outdir, export)
        if not os.path.exists(path):
            print(f"{RED}Error{RESET}: {path} does not exist")
            success = False

    return success


def main():
    parser = argparse.ArgumentParser(description="ostree image tests")
    parser.add_argument(
        "--store",
        metavar="DIRECTORY",
        type=os.path.abspath,
        default=None,
        help="directory where intermediary os trees are stored")
    parser.add_argument(
        "--output-directory",
        metavar="DIRECTORY",
        type=os.path.abspath,
        default=None,
        help="directory where result objects are stored")
    parser.add_argument(
        "--manifest",
        metavar="FILE",
        type=str,
        required=True,
        help="manifest to build")
    parser.add_argument(
        "--export",
        metavar="ID",
        type=str,
        nargs="+",
        required=True,
        help="expected export filepaths (can be passed multiple times)")
    args = parser.parse_args()

    print(f"Running in {os.path.realpath(os.curdir)}")

    tmpdir = "/var/osbuild/tmp"
    os.makedirs(tmpdir, exist_ok=True)
    with tempfile.TemporaryDirectory(dir=tmpdir) as tmp:
        success = run_tests(args, tmp)

    if not success:
        sys.exit(1)


if __name__ == "__main__":
    main()