#!/usr/bin/env python3
"""Create SYSTEM on ARM from a mounted prepared root, burn it, then boot the result."""
import io
import json
from pathlib import Path
import shlex
import subprocess
import sys
import tarfile
import tempfile
import time

project=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(project/'tools'))
from vm_test import VM
from image_formats import digest
work=Path(tempfile.mkdtemp(prefix='m9-system.',dir=project/'out'))
# Reuse only the independently observed virtual hardware map; all images below
# are rebuilt from the current packaged rootfs. No workstation disk is exposed.
runs=sorted((project/'out').glob('m9-vm.*'),key=lambda p:p.stat().st_mtime,reverse=True)
previous=next(p for p in runs if (p/'verified-data.ext4').is_file())
with tarfile.open(previous/'burn-system.tar') as archive:
    bays=archive.extractfile('etc/fds/bays.toml').read()
assert b'3 = 3' in bays,'Run the M9 source-DATA protection test first'

def build(name,rootfs):
    output=work/name;output.mkdir()
    with (work/f'{name}.image.log').open('wb') as log:
        subprocess.run([str(project/'image/build-system-cartridge'),'--rootfs',str(rootfs),'--output-directory',str(output)],check=True,stdout=log,stderr=subprocess.STDOUT)
    return output/'system.img'
source=build('prepared-source',project/'out/rootfs-cli.tar')
replacements={'etc/fds/bays.toml':bays,'usr/libexec/fds/console-session':b'#!/bin/bash\n# Isolated test image only.\nexec env HOME=/root bash --login\n'}
replacements.update({f'usr/bin/{name}':(project/'out'/name).read_bytes() for name in ['fds','fds-burn','fds-cartridged','fds-profile','fds-inspect','fds-eject']})
with tarfile.open(project/'out/rootfs-cli.tar') as src,tarfile.open(work/'creator.tar','w',format=tarfile.PAX_FORMAT) as out:
    seen=set()
    for member in src:
        if member.name in replacements:seen.add(member.name);continue
        out.addfile(member,src.extractfile(member) if member.isfile() else None)
    assert seen==replacements.keys()
    for path,data in replacements.items():
        member=tarfile.TarInfo(path);member.size=len(data);member.mode=0o755 if path.startswith('usr/') else 0o644
        out.addfile(member,io.BytesIO(data))
creator=build('creator',work/'creator.tar')
target=work/'created-system.img'
with target.open('wb') as stream:stream.truncate(640*1024*1024)

def capture(vm,command,timeout=300):
    vm.send('printf "\\nM9_SYSTEM_BEGIN\\n"; '+command+'; printf "\\nM9_SYSTEM_END\\n"')
    return vm.expect(rb'^M9_SYSTEM_BEGIN\r?\n(.*?)\r?\nM9_SYSTEM_END\r?$',timeout=timeout).group(1).decode().strip()
def user(vm,args):
    output=capture(vm,'s6-setuidgid fds fds --json '+shlex.join(args)+' 2>/home/fds/system-command.err; printf "\\nSTATUS:%s" "$?"')
    text,status=output.rsplit('STATUS:',1)
    assert status.strip()=='0',(args,output,capture(vm,'cat /home/fds/system-command.err; tail -n 40 /run/log/cartridged/current'))
    return json.loads(text)

with VM(work,'create',creator,system_usb=True,extra=['-m','2048','-drive',f'file={target},if=none,id=target,format=raw','-device','usb-storage,id=targetusb,drive=target,bus=xhci.0,port=2,serial=M9SYSTEM']) as vm:
    vm.expect(rb'FDS# ')
    capture(vm,'fds-boottrace mark console-ready; cd /home/fds')
    vm.qmp('blockdev-add',{'driver':'raw','node-name':'preparedsource','read-only':True,'file':{'driver':'file','filename':str(source)}})
    vm.qmp('device_add',{'driver':'usb-storage','id':'sourceusb','drive':'preparedsource','bus':'xhci.0','port':'3','serial':'M9PREPARED'})
    deadline=time.monotonic()+45
    while True:
        bay=user(vm,['bay','3'])['bays'][0]
        if bay['state']=='mounted_read_only':break
        assert time.monotonic()<deadline,bay
    assert bay['manifest']['cartridge']['class']=='system'
    result=capture(vm,'fds-burn create system /run/fds/media/03 /home/fds/new-system.img >/home/fds/system-info.json 2>/home/fds/create-system.log; printf "CREATE_STATUS:%s" "$?"')
    assert result=='CREATE_STATUS:0',(result,capture(vm,'tail -n 40 /home/fds/create-system.log'))
    info=json.loads(capture(vm,'cat /home/fds/system-info.json'))
    assert info['class']=='system' and info['bytes']<target.stat().st_size
    (work/'created-source-info.json').write_text(json.dumps(info,indent=2)+'\n')
    print('PASS: actual ARM SYSTEM construction from a complete mounted prepared root',flush=True)
    prepared=user(vm,['burn','system','/home/fds/new-system.img','BAY02'])
    assert prepared['image_sha256']==info['sha256']
    result=user(vm,['burn','confirm',prepared['id'],prepared['confirmation']])
    assert result['phase']=='complete'
    (work/'write-result.json').write_text(json.dumps(result,indent=2)+'\n')
    user(vm,['eject','BAY02'])
    print('PASS: ordinary-user confirmed SYSTEM write, device readback and SAFE eject',flush=True)

layout=json.loads(subprocess.check_output(['sfdisk','--json',str(target)]))['partitiontable']
assert len(layout['partitions'])==1 and layout['partitions'][0]['name']=='FDS_SYSTEM'
assert layout['lastlba']==target.stat().st_size//512-34
(work/'layout.json').write_text(json.dumps(layout,indent=2)+'\n')
(work/'target.sha256').write_text(digest(target)+'  created-system.img\n')
with VM(work,'boot-written-system',target,system_usb=True) as vm:
    vm.expect(rb'FDS> ',timeout=120)
    assert capture(vm,'id -u')=='1000'
    identity=json.loads(capture(vm,'fds --json info'))
    assert identity['target']=='aarch64 static-musl' and identity['pid1'].endswith('s6-svscan'),identity
    assert capture(vm,'stat -c "%u:%g:%a" /usr/lib/utempter/utempter')=='0:14:2711'
    assert capture(vm,'touch /etc/m9-write-rejection 2>/dev/null; printf "STATUS:%s" "$?"')=='STATUS:1'
    assert capture(vm,'pgrep -x dasungd').isdigit()
    assert 'Terminus' in capture(vm,'fc-match Terminus')
    assert capture(vm,'test ! -d /home/fds/.cache/fontconfig; printf "CACHE_STATUS:%s" "$?"')=='CACHE_STATUS:0'
    (work/'booted-identity.json').write_text(json.dumps(identity,indent=2)+'\n')
    print('PASS: the SYSTEM created and written on ARM boots through stage0 to native s6 and the ordinary FDS console',flush=True)
print(f'PASS: M9 SYSTEM creation/write/boot evidence: {work}')
print('SKIP: physical Raspberry Pi boot, storage power-loss behavior and throughput')