#!/usr/bin/env python3

import argparse
import subprocess
import json
import os
import sys
import shutil
import yaml

# Just test these boot to start with.  In the future we should at least
# do ostree upgrades with uefi etc.  But we don't really need the *full*
# suite...if podman somehow broke with nvme or uefi I'd be amazed and impressed.
BASIC_SCENARIOS = ["nvme=true", "firmware=uefi", "firmware=uefi-secure"]

cosa_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, cosa_dir)

from cosalib import cmdlib
from cosalib.builds import Builds

basearch = cmdlib.get_basearch()

# Parse args and dispatch
parser = argparse.ArgumentParser()
parser.add_argument("--build", help="Build ID")
parser.add_argument("--basic-qemu-scenarios", help="Run the basic test across uefi-secure,nvme etc.", action='store_true')
parser.add_argument("--output-dir", help="Output directory")
parser.add_argument("subargs", help="Remaining arguments for kola", nargs='*',
                    default=['run'])
args = parser.parse_args()

builds = Builds()
if args.build is None:
    args.build = builds.get_latest()
builddir = builds.get_build_dir(args.build)
with open(os.path.join(builddir, "meta.json")) as f:
    buildmeta = json.load(f)

# automatically add blacklisted tests specified in the src config
blacklist_args = []
blacklist_path = "src/config/kola-blacklist.yaml"
if os.path.isfile(blacklist_path):
    with open(blacklist_path) as f:
        blacklist = yaml.safe_load(f)
        for obj in (blacklist or []):
            # if there are any arches specified, skip tests only for those arches. if not skip unconditionally.
            if basearch in obj.get('arches', [basearch]):
                print(f"⚠️  Skipping kola test pattern \"{obj['pattern']}\":")
                print(f"⚠️  {obj['tracker']}")
                blacklist_args += ['--blacklist-test', obj['pattern']]

qemuimg = buildmeta['images'].get('qemu')
if qemuimg is None:
    raise SystemExit(f"No qemu image in build: {args.build}")
qemupath = os.path.join(builddir, qemuimg['path'])

# XXX: teach to kola to auto-detect based on prefix; see discussions in
# https://github.com/coreos/coreos-assembler/pull/85
kolaargs = ['kola']
bn = os.path.basename(qemupath)
if not any([x in args.subargs for x in ["-b", "--distro"]]):
    if bn.startswith("rhcos-"):
        kolaargs.extend(['-b', 'rhcos'])
    else:
        kolaargs.extend(['-b', 'fcos'])

print(f"qemu path: {qemupath}")
ignition_version = cmdlib.disk_ignition_version(qemupath)
print(f"Using ignition version {ignition_version}")

if ignition_version == "2.2.0":
    kolaargs.extend(["--ignition-version", "v2"])

if os.getuid() != 0 and not ('-p' in args.subargs):
    kolaargs.extend(['-p', 'qemu-unpriv'])

# tpm devices for qemu not supported on s390x/ppc
if basearch in ["s390x", "ppc64le"]:
    kolaargs.extend(['--qemu-swtpm=false'])

# shellcheck disable=SC2086
kolaargs.extend(['--qemu-image', qemupath])
outputdir = args.output_dir or "tmp/kola"
kolaargs.extend(['--output-dir', outputdir])
kolaargs.extend(args.subargs)

kolaargs.extend(blacklist_args)

env = dict(os.environ)
# By default, store ssh agent in tmp/ too so it can be
# conveniently found.
if args.output_dir is None:
    kola_ssh_dir = 'tmp/kola-ssh'
    if os.path.isdir(kola_ssh_dir):
        shutil.rmtree(kola_ssh_dir)
    os.mkdir(kola_ssh_dir)
    env['MANTLE_SSH_DIR'] = kola_ssh_dir

if args.basic_qemu_scenarios:
    for scenario in BASIC_SCENARIOS:
        subargs = kolaargs + ['--qemu-' + scenario, 'basic']
        print(subprocess.list2cmdline(subargs), flush=True)
        subprocess.check_call(subargs)
else:
    # flush before exec; see https://docs.python.org/3.7/library/os.html#os.execvpe
    print(subprocess.list2cmdline(kolaargs), flush=True)
    os.execvpe('kola', kolaargs, env)
