#!/usr/bin/env python3
"""Confirmed writes by the ordinary FDS user to disposable virtual USB media."""
import hashlib
import io
import itertools
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, gpt, LINUX_FILESYSTEM
work = Path(tempfile.mkdtemp(prefix='m9-vm.', dir=project/'out'))
controller = 'qemu-xhci,id=xhci,addr=05.0'

# A complete legacy PROGRAM compatibility image generated on the workstation.
image_runs = sorted((project/'out').glob('m9-images.*'), key=lambda p:p.stat().st_mtime, reverse=True)
program = next(p/'root/m9-test/program.img' for p in image_runs if (p/'program.json').is_file())
assert program.is_file(), 'Run make media-image-test first'

# Bypass the native creator's manifest validation to exercise hostile input at
# the worker boundary. The filesystem and GPT themselves are valid.
malformed_root=work/'malformed-source'
(malformed_root/'FDS').mkdir(parents=True)
(malformed_root/'FDS/CARTRIDGE.TOML').write_text('format = 1\n[cartridge]\nname = "'+'x'*(60*1024))
malformed_fs=work/'malformed.erofs'
subprocess.run([str(project/'tools/in-image-tools'),'mkfs.erofs','--quiet','-b','4096','-T','0','-L','FDS_PROGRAM',str(malformed_fs),str(malformed_root)],check=True)
malformed_image=work/'malformed.img'
gpt(malformed_image,[('FDS_PROGRAM',LINUX_FILESYSTEM,malformed_fs)])

def fixture(name, hub=None):
    replacements = {'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']})
    if hub:
        replacements['etc/fds/bays.toml'] = f'[front]\nhub = "{hub}"\n[front.ports]\n1 = 1\n2 = 2\n3 = 3\n'.encode()
    extra = {'usr/share/fds/m9-program.img': program.read_bytes(),
             'usr/share/fds/m9-malformed.img': malformed_image.read_bytes()}
    archive = work/f'{name}.tar'
    with tarfile.open(project/'out/rootfs-cli.tar') as src, tarfile.open(archive,'w',format=tarfile.PAX_FORMAT) as out:
        assert src.getmember('usr/bin/fds-burn'), 'Build the M9 rootfs first'
        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,**extra}.items():
            member=tarfile.TarInfo(path);member.size=len(data)
            member.mode=0o755 if path.endswith('console-session') or path.startswith('usr/bin/') else 0o644
            out.addfile(member,io.BytesIO(data))
    destination=work/name;destination.mkdir()
    with (work/f'{name}.image.log').open('wb') as log:
        subprocess.run([str(project/'image/build-system-cartridge'),'--rootfs',str(archive),'--output-directory',str(destination)],check=True,stdout=log,stderr=subprocess.STDOUT)
    return destination/'system.img'

def capture(vm, command):
    vm.send('printf "\\nM9_BEGIN\\n"; '+command+'; printf "\\nM9_END\\n"')
    return vm.expect(rb'^M9_BEGIN\r?\n(.*?)\r?\nM9_END\r?$',timeout=180).group(1).decode().strip()

def user(vm,args,ok=True):
    command='s6-setuidgid fds fds --json '+shlex.join(args)+' 2>/home/fds/command.err; printf "\\nSTATUS:%s" "$?"'
    output=capture(vm,command)
    text,code=output.rsplit('STATUS:',1)
    assert (code.strip()=='0')==ok,(args,output,capture(vm,'cat /home/fds/command.err; tail -n 30 /run/log/cartridged/current'))
    return json.loads(text) if ok else capture(vm,'cat /home/fds/command.err')

def wait_bay(vm,state,number=2):
    deadline=time.monotonic()+30
    while True:
        value=user(vm,['bay',str(number)])['bays'][0]
        if value['state']==state:return value
        if time.monotonic()>deadline:raise AssertionError(value)

def restart(vm,crash=False):
    stop='s6-svc -O /run/service/cartridged && s6-svc -k' if crash else 's6-svc -d'
    output=capture(vm,stop+' /run/service/cartridged && s6-svwait -d -t 15000 /run/service/cartridged && s6-svc -u /run/service/cartridged && s6-svwait -U -t 15000 /run/service/cartridged && printf RESTARTED')
    assert 'RESTARTED' in output,output

sequence=itertools.count()
def attach(vm,path=None,port=2):
    path=path or blank
    node=f'media{next(sequence)}'
    vm.qmp('blockdev-add',{'driver':'raw','node-name':node,'file':{'driver':'file','filename':str(path)}})
    vm.qmp('device_add',{'driver':'usb-storage','id':'targetusb' if port==2 else 'sourceusb','drive':node,'bus':'xhci.0','port':str(port),'serial':'M9TARGET' if port==2 else 'M9SOURCE'})

blank=work/'cartridge.img'
with blank.open('wb') as stream:stream.truncate(96*1024*1024)
unchanged=digest(blank)
probe=fixture('probe')
with VM(work,'probe',probe,system_usb=True,extra=['-drive',f'file={blank},if=none,id=target,format=raw','-device','usb-storage,id=targetusb,drive=target,bus=xhci.0,port=2,serial=M9TARGET']) as vm:
    vm.expect(rb'FDS# ')
    topology=json.loads(capture(vm,'fds --json topology'))
    hub=next(d['topology'] for d in topology['unmapped'] if d.get('serial')=='M9TARGET').rsplit('/',1)[0]

system=fixture('burn-system',hub)
with VM(work,'burn',system,system_usb=True,extra=['-drive',f'file={blank},if=none,id=target,format=raw','-device','usb-storage,id=targetusb,drive=target,bus=xhci.0,port=2,serial=M9TARGET']) as vm:
    vm.expect(rb'FDS# ')
    capture(vm,'fds-boottrace mark console-ready; cd /home/fds')
    disk=user(vm,['inspect','BAY02']);assert disk['bytes']==blank.stat().st_size and disk['protected'] is None
    assert user(vm,['inspect','BAY01'])['protected']
    user(vm,['burn','program','/usr/share/fds/m9-program.img','BAY01'],ok=False)
    user(vm,['burn','system','/usr/share/fds/m9-program.img','BAY02'],ok=False)
    assert digest(blank)==unchanged
    daemon_pid=capture(vm,'s6-svstat -o pid /run/service/cartridged')
    error=user(vm,['burn','program','/usr/share/fds/m9-malformed.img','BAY02'],ok=False)
    assert 'Invalid cartridge manifest' in error and len(error)<2200,error
    assert not any(ord(c)<32 for c in error),repr(error)
    assert capture(vm,'s6-svstat -o pid /run/service/cartridged')==daemon_pid
    assert len(user(vm,['bays'])['bays'])==12 and digest(blank)==unchanged
    print('PASS: oversized manifest diagnostics remain bounded; daemon stays responsive and target unchanged',flush=True)
    prepared=user(vm,['burn','program','/usr/share/fds/m9-program.img','BAY02'])
    assert prepared['phase']=='awaiting_confirmation' and prepared['image_sha256']==digest(program)
    user(vm,['burn','confirm',prepared['id'],'ERASE BAY02 wrong-operation'],ok=False)
    assert user(vm,['burn','status',prepared['id']])['phase']=='awaiting_confirmation'
    assert digest(blank)==unchanged,'Preview or rejected confirmation modified the disk'
    user(vm,['burn','cancel',prepared['id']],ok=False)
    assert digest(blank)==unchanged
    print('PASS: active SYSTEM protection, class validation, bound confirmation, preview and cancellation do not write',flush=True)

    capture(vm,'cp /usr/share/fds/m9-program.img /home/fds/mutable.img; chown fds:fds /home/fds/mutable.img')
    prepared=user(vm,['burn','program','/home/fds/mutable.img','BAY02'])
    capture(vm,"printf changed | dd of=/home/fds/mutable.img bs=1 seek=2097152 conv=notrunc status=none")
    error=user(vm,['burn','confirm',prepared['id'],prepared['confirmation']],ok=False)
    assert 'changed after confirmation' in error,error
    assert digest(blank)==unchanged
    print('PASS: source mutation after preview is rejected before the target is written',flush=True)

    prepared=user(vm,['burn','program','/usr/share/fds/m9-program.img','BAY02'])
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    attach(vm)
    # The old worker holds the old insertion's descriptor; it must never follow
    # the new kernel disk in the same physical port.
    deadline=time.monotonic()+30
    while True:
        observed=capture(vm,'s6-setuidgid fds fds --json inspect BAY02 2>/home/fds/enumeration.err')
        if observed.startswith('{') and json.loads(observed)['diskseq']!=disk['diskseq']:break
        assert time.monotonic()<deadline,observed
    error=user(vm,['burn','confirm',prepared['id'],prepared['confirmation']],ok=False)
    assert 'identity changed' in error,error
    assert digest(blank)==unchanged
    print('PASS: replacement in the same bay invalidates the old confirmation without writing',flush=True)

    prepared=user(vm,['burn','program','/usr/share/fds/m9-program.img','BAY02'])
    result=user(vm,['burn','confirm',prepared['id'],prepared['confirmation']])
    assert result['phase']=='complete'
    assert wait_bay(vm,'safe')['mount'] is None
    restart(vm);assert wait_bay(vm,'safe')['mount'] is None
    # Removing and reinserting the same bytes is a new kernel diskseq. It can be mounted normally.
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    attach(vm)
    cartridge=wait_bay(vm,'mounted_read_only');assert cartridge['manifest']['cartridge']['class']=='program'
    assert user(vm,['inspect','BAY02'])['diskseq']!=disk['diskseq']
    user(vm,['burn','program','/usr/share/fds/m9-program.img','BAY02'],ok=False)
    user(vm,['eject','BAY02']);wait_bay(vm,'safe')
    print('PASS: verified PROGRAM write, larger-disk backup GPT, SAFE across restart, reinsertion and mounted-target refusal',flush=True)

    capture(vm,'mkdir -p /home/fds/app/bin; cp /usr/bin/fds /home/fds/app/bin/report; chown -R fds:fds /home/fds/app')
    error=user(vm,['format','program','/home/fds/app','BAY02','--label','M9 TOOLS','--id','fds.m9tools'],ok=False)
    assert 'workstation' in error.lower(),error
    # Existing PROGRAM images remain launchable; newly built software images
    # are covered by the native workstation/emulator integration suite.
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty');attach(vm)
    cartridge=wait_bay(vm,'mounted_read_only');assert cartridge['manifest']['cartridge']['id']=='fds.test.program'
    assert user(vm,['run','2','--','fds','info'])['started_pid']>1
    user(vm,['eject','BAY02'])
    print('PASS: workstation-only PROGRAM creation and explicit legacy application launch',flush=True)

    prepared=user(vm,['format','data','BAY02','--label','M9 DATA','--size-mib','32'])
    assert prepared['phase']=='awaiting_confirmation'
    # Leave an actual write running while querying the ordinary control endpoint.
    command='s6-setuidgid fds fds --json burn confirm '+shlex.join([prepared['id'],prepared['confirmation']])+' >/home/fds/write-result.json 2>/home/fds/write-result.err &'
    capture(vm,'('+command+')')
    assert len(user(vm,['bays'])['bays'])==12
    result=user(vm,['burn','wait',prepared['id']]);assert result['phase']=='complete'
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    attach(vm)
    cartridge=wait_bay(vm,'mounted_read_write');assert cartridge['manifest']['cartridge']['name']=='M9 DATA'
    assert 'WROTE_DATA' in capture(vm,"s6-setuidgid fds /bin/bash -c 'printf persistent > /data/test-write' && printf WROTE_DATA")
    user(vm,['eject','BAY02']);wait_bay(vm,'safe')
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    # QEMU flushes the backend on safe removal; retain DATA for an independent host fsck.
    saved=work/'verified-data.img'
    import shutil
    shutil.copyfile(blank,saved)
    attach(vm)
    wait_bay(vm,'mounted_read_write');user(vm,['eject','BAY02'])
    print('PASS: on-target DATA formatting, responsive status during writing, unprivileged use and safe eject',flush=True)

    # Source DATA remains protected while an image descriptor is held in the
    # writer's private mount namespace, even though the destination is another bay.
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    attach(vm,port=3);wait_bay(vm,'mounted_read_write',3)
    other=work/'other-target.img'
    with other.open('wb') as stream:stream.truncate(96*1024*1024)
    attach(vm,other)
    deadline=time.monotonic()+30
    while True:
        observed=capture(vm,'fds --json inspect BAY02 2>/home/fds/enumeration.err')
        if observed.startswith('{'):break
        assert time.monotonic()<deadline
    capture(vm,'s6-setuidgid fds cp /usr/share/fds/m9-program.img /data/source.img')
    prepared=user(vm,['burn','program','/data/source.img','BAY02'])
    assert 'media operation is active' in user(vm,['eject','3'],ok=False)
    assert user(vm,['bay','3'])['bays'][0]['mount']=='/data'
    user(vm,['burn','cancel',prepared['id']],ok=False)
    user(vm,['eject','3']);wait_bay(vm,'safe',3)
    vm.qmp('device_del',{'id':'sourceusb'});wait_bay(vm,'empty',3)
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    attach(vm);wait_bay(vm,'mounted_read_write');user(vm,['eject','2'])
    print('PASS: a source DATA cartridge cannot be ejected while its image is held by the writer',flush=True)

    prepared=user(vm,['format','environment','BAY02','--profile','cli','--label','CONSOLE ENVIRONMENT'])
    assert user(vm,['burn','confirm',prepared['id'],prepared['confirmation']])['phase']=='complete'
    vm.qmp('device_del',{'id':'targetusb'});wait_bay(vm,'empty')
    attach(vm)
    cartridge=wait_bay(vm,'mounted_read_only');assert cartridge['manifest']['activation']['profile']=='cli'
    user(vm,['eject','BAY02'])
    prepared=user(vm,['burn','program','/usr/share/fds/m9-program.img','BAY02'])
    restart(vm,crash=True)
    assert 'restarted' in user(vm,['burn','status',prepared['id']],ok=False)
    assert wait_bay(vm,'failed')['mount'] is None
    print('PASS: on-target ENVIRONMENT formatting and fail-closed interrupted-operation restart',flush=True)

layout=json.loads(subprocess.check_output(['sfdisk','--json',str(saved)]))['partitiontable']['partitions'][0]
assert layout['name']=='FDS_DATA'
payload=work/'verified-data.ext4'
with saved.open('rb') as src,payload.open('wb') as dst:
    src.seek(layout['start']*512);remaining=layout['size']*512
    while remaining:
        data=src.read(min(1024*1024,remaining));assert data;dst.write(data);remaining-=len(data)
subprocess.run(['e2fsck','-fn',str(payload)],check=True,stdout=subprocess.DEVNULL)
assert subprocess.check_output(['debugfs','-R','cat /test-write',str(payload)],stderr=subprocess.DEVNULL)==b'persistent'
print(f'PASS: M9 confirmed-write integration: {work}')
print('SKIP: physical media, power-loss behavior and Pi write throughput')