#!/usr/bin/env python3
"""Real write, flush, readback, cancellation and worker-crash failure boundaries."""
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 gpt, digest, LINUX_FILESYSTEM
work=Path(tempfile.mkdtemp(prefix='m11-faults.',dir=project/'out'))
inputs=[Path(__file__),project/'tools/vm_test.py',project/'out/rootfs-cli.tar',
        project/'out/kernel/boot/kernel_2712.img',project/'out/fds-initramfs.img']
(work/'inputs.sha256').write_text(''.join(f'{digest(path)}  {path.relative_to(project)}\n' for path in inputs))
(work/'qemu-version.txt').write_bytes(subprocess.check_output([str(project/'tools/in-void'),'qemu-system-aarch64','--version']))
controller='qemu-xhci,id=xhci,addr=05.0'
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())
# A valid padded PROGRAM filesystem makes the transfer long enough to observe
# cancellation after real writes. Its contents are still independently checked.
layout=json.loads(subprocess.check_output(['sfdisk','--json',str(program)]))['partitiontable']['partitions'][0]
large_fs=work/'large.erofs'
with program.open('rb') as src,large_fs.open('xb') as dst:
    src.seek(layout['start']*512);dst.write(src.read(layout['size']*512));dst.truncate(256*1024*1024)
large=work/'large.img';gpt(large,[('FDS_PROGRAM',LINUX_FILESYSTEM,large_fs)])
subprocess.run([str(project/'tools/in-image-tools'),'fsck.erofs',str(large_fs)],check=True)

def query(vm,args):return json.loads(vm.capture('s6-setuidgid fds fds --json '+shlex.join(args)))
def wait(vm,read,predicate,timeout=90):
    deadline=time.monotonic()+timeout
    while True:
        value=read()
        if predicate(value):return value
        assert time.monotonic()<deadline,value

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'}
    if hub:
        replacements['etc/fds/bays.toml']=f'[usb]\nhub="{hub}"\n[usb.ports]\n2=2\n'.encode()
    additions={'usr/share/fds/m11-small.img':program.read_bytes(),'usr/share/fds/m11-large.img':large.read_bytes()}
    with tarfile.open(project/'out/rootfs-cli.tar') as src,tarfile.open(work/f'{name}.tar','w',format=tarfile.PAX_FORMAT) as dst:
        seen=set()
        for member in src:
            if member.name in replacements:seen.add(member.name);continue
            dst.addfile(member,src.extractfile(member) if member.isfile() else None)
        assert seen==replacements.keys()
        for path,data in {**replacements,**additions}.items():
            member=tarfile.TarInfo(path);member.size=len(data);member.mode=0o755 if path.endswith('console-session') else 0o644
            dst.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(work/f'{name}.tar'),'--output-directory',str(destination)],check=True,stdout=log,stderr=subprocess.STDOUT)
    return destination/'system.img'

probe=fixture('probe')
with VM(work,'probe',probe,extra=['-device',controller,'-device','usb-kbd,bus=xhci.0,port=2']) as vm:
    vm.expect(rb'FDS# ')
    report=wait(vm,lambda:query(vm,['topology']),lambda r:bool(r['unmapped']))
    hub=next(d['topology'] for d in report['unmapped'] if '03' in d['interfaces']).rsplit('/',1)[0]
# USB storage negotiates SuperSpeed and has a distinct root-protocol identity.
system=fixture('mapped',hub.replace(':usb2',':usb3'))

def attach(vm,disk,rule=None,throttle=False):
    node={'driver':'raw','node-name':'raw','file':{'driver':'file','filename':str(disk)}}
    vm.qmp('blockdev-add',node)
    top='raw'
    if rule:
        vm.qmp('blockdev-add',{'driver':'blkdebug','node-name':'fault','image':top,'inject-error':[{'event':'none','errno':5,'once':False,**rule}]})
        top='fault'
    if throttle:
        vm.qmp('object-add',{'qom-type':'throttle-group','id':'limit','limits':{'bps-write':8*1024*1024}})
        vm.qmp('blockdev-add',{'driver':'throttle','node-name':'limited','throttle-group':'limit','file':top})
        top='limited'
    vm.qmp('device_add',{'driver':'usb-storage','id':'target','drive':top,'bus':'xhci.0','port':'2','serial':'M11-FAULT'})
    wait(vm,lambda:vm.capture('fds --json inspect BAY02 2>/tmp/m11-enumeration.err; true'),lambda s:s.startswith('{') and json.loads(s)['bytes']==disk.stat().st_size)

def job_record(vm):return json.loads(vm.capture('cat /run/fds/burn/02.json'))['job']
def failed(vm,job):
    observed=wait(vm,lambda:job_record(vm),lambda j:j['phase']=='failed')
    assert observed['id']==job['id'] and observed['error'],observed
    assert query(vm,['bay','2'])['bays'][0]['state']!='safe'
    wait(vm,lambda:vm.capture('pgrep -x fds-burn || test "$?" = 1'),lambda s:not s)
    vm.capture('s6-svc -d /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')
    assert job_record(vm)['phase']=='failed'
    assert query(vm,['bay','2'])['bays'][0]['state']=='failed'
    vm.capture('test ! -e /run/fds/ejected/02')
    (work/f'{vm.name}-result.json').write_text(json.dumps(observed,indent=2)+'\n')
    (work/f'{vm.name}-kernel.log').write_text(vm.capture('dmesg')+'\n')
    return observed

# Errors affect sectors beyond the GPT checks, so preview remains non-destructive.
for name,rule in [('write-error',{'iotype':'write','sector':2056}),
                  ('flush-error',{'iotype':'flush'}),
                  ('readback-error',{'iotype':'read','sector':2056})]:
    disk=work/f'{name}.img'
    with disk.open('xb') as f:f.truncate(256*1024*1024)
    before=digest(disk)
    with VM(work,name,system,extra=['-device',controller]) as vm:
        vm.expect(rb'FDS# ');attach(vm,disk,rule)
        job=query(vm,['burn','program','/usr/share/fds/m11-small.img','BAY02'])
        assert job['phase']=='awaiting_confirmation' and digest(disk)==before
        error=vm.capture('s6-setuidgid fds fds burn confirm '+shlex.join([job['id'],job['confirmation']]),ok=False)
        result=failed(vm,job)
        assert 'os error 5' in result['error'],(name,error,result)
    assert digest(disk)!=before,(name,'No write was observed')
    print(f'PASS: {name}: actual EIO after confirmation; failed across restart, never SAFE',flush=True)

for name in ['cancel-writing','kill-worker','kill-daemon']:
    disk=work/f'{name}.img'
    with disk.open('xb') as f:f.truncate(384*1024*1024)
    before=digest(disk)
    with VM(work,name,system,extra=['-device',controller]) as vm:
        vm.expect(rb'FDS# ');attach(vm,disk,throttle=True)
        job=query(vm,['burn','program','/usr/share/fds/m11-large.img','BAY02'])
        vm.capture('(s6-setuidgid fds fds burn confirm '+shlex.join([job['id'],job['confirmation']])+' >/tmp/m11-confirm.out 2>&1 & )')
        progress=wait(vm,lambda:job_record(vm),lambda j:j['phase']=='writing' and j['progress_bytes']>=64*1024*1024,timeout=180)
        assert progress['progress_bytes']<job['image_bytes'],progress
        observed=wait(vm,lambda:vm.qmp('query-blockstats'),
            lambda rows:any('/target/' in r.get('qdev','') and r['stats']['wr_bytes']>=64*1024*1024 for r in rows),timeout=180)
        progress=job_record(vm)
        assert progress['phase']=='writing' and progress['progress_bytes']<job['image_bytes'],progress
        (work/f'{name}-interruption.json').write_text(json.dumps({'job':progress,'blockstats':observed},indent=2)+'\n')
        if name=='cancel-writing':
            vm.capture('s6-setuidgid fds fds burn cancel '+job['id'],ok=False)
        elif name=='kill-worker':
            vm.capture('pkill -KILL -x fds-burn')
        else:
            vm.capture('s6-svc -O /run/service/cartridged && s6-svc -k /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')
        result=failed(vm,job)
        assert 'cancelled' in result['error'] if name=='cancel-writing' else ('exited' in result['error'] or 'restarted' in result['error']),result
        assert len(query(vm,['bays'])['bays'])==12
    assert digest(disk)!=before,'Cancellation/crash occurred before any physical write'
    print(f'PASS: {name}: interrupted after at least 64 MiB, service remains usable, no false SAFE',flush=True)
link=project/'out/m11-faults-latest';link.unlink(missing_ok=True);link.symlink_to(work.name)
print(f'PASS: M11 media failure evidence: {work}',flush=True)
print('SKIP: flash-controller behavior, physical removal, power cuts and Pi throughput')