#!/usr/bin/env python3
"""Continue a successfully prepared installed VM through faults/apply/cold boot.

Run as one durable systemd service with CPUQuota=200%, MemoryMax=8G,
MemorySwapMax=0. Each child takes the shared heavy lock inside that job. Faults
operate on independent disk copies, never restore the active VM's financial data.
"""
import argparse
import json
from pathlib import Path
import subprocess
import sys
import time


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--prepared-report', type=Path, required=True)
    parser.add_argument('--baseline-boot', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--wait-prepared', type=int, default=0)
    parser.add_argument('--preparation-faults-only', action='store_true',
                        help='run offline/interrupted preparation only, even after failed prepare')
    parser.add_argument('--preparation-fault', action='append',
                        choices=['offline','interrupted-before','bad-hash','space-prepare','partial-dpkg'])
    parser.add_argument('--fault',action='append',choices=['locked-apt','space','partial-dpkg','offline','missing-dependency','interrupted-before'])
    args = parser.parse_args()
    if args.preparation_fault and not args.preparation_faults_only:
        parser.error('--preparation-fault requires --preparation-faults-only')
    if args.fault and args.preparation_faults_only:
        parser.error('--fault is for the apply/coldboot matrix')
    deadline = time.monotonic()+args.wait_prepared
    while not (args.prepared_report/'result.json').exists() and time.monotonic()<deadline:
        time.sleep(2)
    prepared = json.loads((args.prepared_report/'result.json').read_text())
    if not prepared['passed'] and not args.preparation_faults_only:
        parser.error('a successful actual prepare receipt is required')
    plan = json.loads((args.prepared_report/'os-ota-stage/plan.json').read_text())
    disk = Path(prepared['disk'])
    args.output.mkdir(mode=0o755, exist_ok=False)
    harness = Path(__file__).with_name('os_ota_installed_vm.py')
    source = Path(__file__).resolve().parents[2]
    old = plan['boot']['running_release']
    new = plan['boot']['expected_release']
    boot = ['--kernel',str(args.baseline_boot/('vmlinuz-'+old)),
            '--initrd',str(args.baseline_boot/('initrd.img-'+old))]
    results = []

    def step(name, *command):
        output = args.output/name
        argv = [sys.executable,'-B',str(harness),'--output',str(output),*map(str,command)]
        with (args.output/(name+'.driver.log')).open('x') as log:
            result = subprocess.run(argv, stdout=log, stderr=subprocess.STDOUT)
        results.append(dict(name=name, exit=result.returncode, command=argv))
        (args.output/'progress.json').write_text(json.dumps(results,indent=2)+'\n')
        return result.returncode == 0

    faults = ('offline','interrupted-before') if args.preparation_faults_only else (
        'locked-apt','space','partial-dpkg','offline','missing-dependency','interrupted-before')
    faults = args.preparation_fault or args.fault or faults
    for fault in faults:
        clone = 'disk-'+fault
        if not step(clone, 'clone','--disk',disk):
            continue
        step('fault-'+fault,'boot','--disk',args.output/clone/'installed.raw',*boot,
             '--guest',source/'tests/regression/os-ota-installed-fault.py','--scenario',fault)
    if not args.preparation_faults_only and step('apply','boot','--disk',disk,*boot,
            '--guest',source/'tests/regression/os-ota-installed-apply.py','--offline'):
        if step('export-updated','export','--disk',disk,'--kernel-release',new):
            step('coldboot','boot','--disk',disk,
                 '--kernel',args.output/'export-updated'/('vmlinuz-'+new),
                 '--initrd',args.output/'export-updated'/('initrd.img-'+new),
                 '--guest',source/'tests/regression/os-ota-installed-postboot.py','--offline','--firmware')
    result = dict(passed=all(x['exit']==0 for x in results) and (
                      args.preparation_faults_only or any(x['name']=='coldboot' for x in results)),
                  preparation_faults_only=args.preparation_faults_only,
                  selected_faults=list(faults),
                  production_approved=False, native_D09_transport=False, results=results)
    (args.output/'result.json').write_text(json.dumps(result,indent=2)+'\n')
    return 0 if result['passed'] else 1


if __name__ == '__main__':
    sys.exit(main())