#!/usr/bin/env python3
"""Ordered native poweroff/reboot, busy refusal and persistent DATA quarantine."""
import base64
import binascii
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 gpt, LINUX_FILESYSTEM, digest
work=Path(tempfile.mkdtemp(prefix='m10-vm.',dir=project/'out'))
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())

def capture(vm,command,ok=True):
    # Preserve diagnostics and return exactly the command's output despite
    # unrelated kernel messages on the same serial device. Never retry a command.
    vm.send('('+command+') >/tmp/m10-response 2>&1; printf "%s" "$?" >/tmp/m10-status; printf "\\nM10_SAVED\\n"')
    vm.expect(rb'^M10_SAVED\r?$',timeout=180)
    for attempt in range(20):
        vm.send('printf "\\nM10_BEGIN\\n"; base64 -w0 /tmp/m10-response; printf "\\n"; sha256sum /tmp/m10-response; cat /tmp/m10-status; printf "\\nM10_END\\n"')
        frame=vm.expect(rb'^M10_BEGIN\r?\n(.*?)\r?\nM10_END\r?$').group(1).decode().replace('\r','')
        try:
            encoded,checksum,status=frame.split('\n')
            payload=base64.b64decode(encoded,validate=True)
            if hashlib.sha256(payload).hexdigest()!=checksum.split()[0]:continue
            assert (status=='0')==ok,(command,status,payload)
            return payload.decode().strip()
        except (ValueError,binascii.Error):continue
    raise AssertionError('Repeated serial transfer corruption')

def query(vm,args):
    return json.loads(capture(vm,'s6-setuidgid fds fds --json '+shlex.join(args)))
def rejected(vm,args):
    return capture(vm,'s6-setuidgid fds fds '+shlex.join(args),ok=False)
def wait(vm,command,condition,timeout=40):
    deadline=time.monotonic()+timeout
    while True:
        value=capture(vm,command)
        if condition(value):return value
        assert time.monotonic()<deadline,(command,value)
def bay(vm,n,state):
    return json.loads(wait(vm,f'fds --json bay {n}',lambda s:json.loads(s)['bays'][0]['state']==state))['bays'][0]
def boot(vm):
    vm.expect(rb'FDS# ')
    capture(vm,'fds-boottrace mark console-ready')

def fixture(name,config=None):
    replacements={'usr/libexec/fds/console-session':b'#!/bin/bash\n# Isolated test image only.\nexec env HOME=/root bash --login\n','etc/fds/xserver':b'xvfb\n'}
    if config:replacements['etc/fds/bays.toml']=config.encode()
    with tarfile.open(project/'out/rootfs-development.tar') as src,tarfile.open(work/f'{name}.tar','w',format=tarfile.PAX_FORMAT) as dst:
        assert src.extractfile('usr/share/fds/image-profile').read().strip()==b'development','Build PROFILE=development first'
        assert src.getmember('usr/bin/fds-power'),'Build the M10 rootfs first'
        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()
        extra={'usr/libexec/fds/m10-writer':(project/'out/m7-writer').read_bytes(),'usr/share/fds/m10-program.img':program.read_bytes()}
        for path,data in {**replacements,**extra}.items():
            member=tarfile.TarInfo(path);member.size=len(data);member.mode=0o755 if path.startswith('usr/libexec/') else 0o644
            dst.addfile(member,io.BytesIO(data))
    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(work/f'{name}.tar'),'--output-directory',str(output)],check=True,stdout=log,stderr=subprocess.STDOUT)
    return output/'system.img'

phases=['frozen','stopping_desktop','stopping_programs','stopping_network','syncing_and_unmounting_data','unmounting_cartridges','prepared','services_stopped']
def finished(vm,reboot=False):
    report=json.loads(vm.expect(rb'^FDS_SHUTDOWN_FINAL (\{[^\r\n]*\})\r?\n',timeout=90).group(1))
    events=report['events'];start=max(i for i,e in enumerate(events) if e['phase']=='frozen')
    final=events[start:]
    assert [e['phase'] for e in final]==phases,report
    assert [e['at_ns'] for e in final]==sorted(e['at_ns'] for e in final)
    ending=vm.expect(rb'\[\s*([0-9]+\.[0-9]+)\] reboot: (Power down|Restarting system)',timeout=30)
    assert (ending.group(2)==b'Restarting system')==reboot
    assert vm.child.wait(timeout=20)==0,'QEMU did not exit after the guest power action'
    report['virtual_total_ms']=float(ending.group(1))*1000-final[0]['at_ns']/1e6
    assert report['virtual_total_ms']>=0
    (work/f'{vm.name}-shutdown.json').write_text(json.dumps(report,indent=2)+'\n')
    print(f'PASS: {vm.name}: ordered native {"reboot" if reboot else "poweroff"}; VM observation {report["virtual_total_ms"]:.3f} ms',flush=True)

probe=fixture('probe-system')
with VM(work,'idle',probe,extra=['-device',controller,'-device','usb-kbd,bus=xhci.0,port=1']) as vm:
    boot(vm)
    topology=query(vm,['topology'])
    hub=next(d['topology'] for d in topology['unmapped'] if '03' in d['interfaces']).rsplit('/',1)[0]
    assert query(vm,['power','status'])['phase']=='idle'
    assert 'requires root' in capture(vm,'s6-setuidgid fds fds-power --shutdown-hook',ok=False)
    assert query(vm,['power','status'])['phase']=='idle'
    vm.send('s6-setuidgid fds fds poweroff')
    finished(vm)

config=''
for name,identity in [('usb2',hub),('usb3',hub.replace(':usb2',':usb3'))]:
    config+=f'[{name}]\nhub={json.dumps(identity)}\n[{name}.ports]\n'+''.join(f'{n}={n}\n' for n in range(1,13))
system=fixture('mapped-system',config)

def data_image(name):
    root=work/f'{name}-files';(root/'FDS').mkdir(parents=True)
    (root/'FDS/CARTRIDGE.TOML').write_text(f'format=1\n[cartridge]\nid="fds.m10.{name}"\nname="M10 DATA"\nclass="data"\nversion="1"\n[media]\nwritable=true\n')
    if name=='ioerror':
        (root/'fault.bin').write_bytes(b'A'*(1024*1024))
        (root/'fault.bin').chmod(0o666)
    fs=work/f'{name}.ext4'
    with fs.open('xb') as f:f.truncate(256*1024*1024)
    subprocess.run(['mke2fs','-q','-t','ext4','-F','-b','4096','-E','root_owner=1000:1000,lazy_itable_init=0,lazy_journal_init=0','-d',str(root),str(fs)],check=True)
    disk=work/f'{name}.img';layout=gpt(disk,[('FDS_DATA',LINUX_FILESYSTEM,fs)])
    return disk,layout

sequence=itertools.count()
def attach(vm,disk,port=2):
    node=f'disk{next(sequence)}'
    vm.qmp('blockdev-add',{'driver':'raw','node-name':node,'file':{'driver':'file','filename':str(disk)}})
    vm.qmp('device_add',{'driver':'usb-storage','id':f'bay{port}','drive':node,'bus':'xhci.0','port':str(port)})
def restart(vm,crash=False):
    stop='s6-svc -O /run/service/cartridged && s6-svc -k' if crash else 's6-svc -d'
    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')
def workload(vm):
    query(vm,['profile','activate','windowmaker'])
    capture(vm,'s6-svwait -U -t 25000 /run/service/desktop-session')
    assert query(vm,['profiles'])['profiles']['ready_ns']
    query(vm,['network','on'])
    assert query(vm,['profiles'])['profiles']['network']
    capture(vm,'rm -f /data/progress')
    assert query(vm,['run','2','--','/usr/libexec/fds/m10-writer'])['started_pid']>1
    wait(vm,'cat /data/progress 2>/dev/null || printf 0',lambda s:s.isdigit() and int(s)>=128)

data,layout=data_image('workload')
with VM(work,'workload-reboot',system,extra=['-device',controller,'-netdev','user,id=net']) as vm:
    boot(vm);attach(vm,data);bay(vm,2,'mounted_read_write')
    capture(vm,'modprobe cdc_ether')
    vm.qmp('device_add',{'driver':'usb-net','id':'ethernet','netdev':'net','bus':'xhci.0','port':'3'})
    wait(vm,'fds --json profiles',lambda s:bool(json.loads(s)['profiles']['network']))
    workload(vm)
    error=capture(vm,'cd /data; s6-setuidgid fds fds poweroff',ok=False)
    assert 'busy' in error.lower(),error
    status=query(vm,['power','status']);assert status['phase']=='blocked' and not status['native_pending']
    capture(vm,"! pgrep -x Xvfb && ! pgrep -x wmaker && ! pgrep -x dhcpcd && grep -q 'populated 0' /sys/fs/cgroup/fds/bay02/cgroup.events && grep -q 'populated 0' /sys/fs/cgroup/fds/network/cgroup.events && pgrep -x dasungd")
    for args in [['run','2','--','/usr/bin/true'],['data','use','2'],['profile','activate','windowmaker'],['network','on'],['rescan']]:
        assert 'frozen' in rejected(vm,args),args
    assert query(vm,['power','resume'])['phase']=='idle'
    error=capture(vm,"s6-setuidgid fds /bin/bash -c 'exec 9>>/data/open-writer; fds poweroff'",ok=False)
    assert 'busy' in error.lower(),error
    assert query(vm,['power','resume'])['phase']=='idle'
    capture(vm,"s6-setuidgid fds /bin/bash -c 'printf persisted > /data/reboot-marker'")
    workload(vm)
    vm.send('s6-setuidgid fds fds reboot')
    finished(vm,reboot=True)
print('PASS: busy working directory/open writer refuse shutdown; frozen operations, explicit resume and actual consumer exit verified',flush=True)

# Inspect the final unmounted filesystem independently; the guest's SAFE marker
# is not accepted as a substitute for filesystem and payload verification.
payload=work/'workload-after.ext4';part=layout['partitions'][0]
with data.open('rb') as src,payload.open('wb') as dst:
    src.seek(part['start']*512);remaining=part['payload_bytes']
    while remaining:
        chunk=src.read(min(1024*1024,remaining));assert chunk;dst.write(chunk);remaining-=len(chunk)
with (work/'workload-fsck.log').open('wb') as log:subprocess.run(['e2fsck','-fn',str(payload)],check=True,stdout=log,stderr=subprocess.STDOUT)
written=work/'stress-after.bin'
subprocess.run(['debugfs','-R',f'dump /stress.bin {written}',str(payload)],check=True,stdout=subprocess.DEVNULL)
assert written.stat().st_size==64*1024*1024
assert digest(written)==hashlib.sha256(bytes(range(256))*(64*1024*1024//256)).hexdigest()
assert subprocess.check_output(['debugfs','-R','cat /reboot-marker',str(payload)],stderr=subprocess.DEVNULL)==b'persisted'
print('PASS: reboot leaves clean ext4 and the exact 64 MiB payload plus ordinary-user files',flush=True)

with VM(work,'native-guard',system,extra=['-device',controller]) as vm:
    boot(vm);attach(vm,data);bay(vm,2,'mounted_read_write')
    assert capture(vm,'cat /data/reboot-marker')=='persisted'
    vm.send('cd /data; /usr/bin/poweroff; printf "\\nNATIVE_REQUESTED\\n"')
    vm.expect(rb'^NATIVE_REQUESTED\r?$')
    status=json.loads(wait(vm,'fds --json power status',lambda s:json.loads(s)['phase']=='blocked'))
    assert status['native_pending'] and 'busy' in status['error'].lower(),status
    assert 'cannot be cancelled' in rejected(vm,['power','resume'])
    vm.send('cd /; s6-setuidgid fds fds poweroff')
    finished(vm)
print('PASS: direct native poweroff remains blocked until DATA can be safely unmounted',flush=True)

blank=work/'blank.img'
with blank.open('wb') as f:f.truncate(96*1024*1024)
unchanged=digest(blank)
with VM(work,'burn-guard',system,extra=['-device',controller]) as vm:
    boot(vm);attach(vm,blank,4)
    bay(vm,4,'unrecognized_storage')
    wait(vm,'fds --json inspect BAY04 2>/tmp/m10-enumeration-error; true',lambda s:s.startswith('{') and json.loads(s)['bytes']==blank.stat().st_size)
    job=query(vm,['burn','program','/usr/share/fds/m10-program.img','BAY04'])
    assert job['phase']=='awaiting_confirmation'
    assert 'active media operation' in rejected(vm,['poweroff'])
    assert 'frozen' in rejected(vm,['burn','confirm',job['id'],job['confirmation']])
    rejected(vm,['burn','cancel',job['id']])
    wait(vm,'pgrep -x fds-burn || test "$?" = 1',lambda s:not s)
    assert query(vm,['bay','4'])['bays'][0]['state']!='safe'
    vm.send('s6-setuidgid fds fds poweroff');finished(vm)
assert digest(blank)==unchanged
print('PASS: active media preparation blocks shutdown; cancellation leaves the target unchanged and never SAFE',flush=True)

quarantine,_=data_image('quarantine')
with VM(work,'quarantine',system,extra=['-device',controller]) as vm:
    boot(vm);attach(vm,quarantine);bay(vm,2,'mounted_read_write')
    capture(vm,"s6-setuidgid fds /bin/bash -c 'printf retained > /data/test'")
    restart(vm,crash=True)
    assert 'interrupted' in bay(vm,2,'error')['detail']
    assert 'Shutdown blocked' in rejected(vm,['poweroff'])
    restart(vm)
    assert query(vm,['power','status'])['phase']=='blocked'
    assert 'quarantined' in rejected(vm,['poweroff'])
    assert query(vm,['bay','2'])['bays'][0]['state']!='safe'
    vm.qmp('device_del',{'id':'bay2'});bay(vm,2,'empty')
    vm.send('s6-setuidgid fds fds poweroff');finished(vm)
print('PASS: DATA quarantine and shutdown freeze survive daemon restart without a false SAFE result',flush=True)

# Inject real EIO into a known allocated file data block, without pulling
# the device or damaging GPT/manifest metadata. Mount-time journal writes do not
# overlap this block, so the failure occurs in an already active DATA session.
fault_disk,fault_layout=data_image('ioerror')
blocks=subprocess.check_output(['debugfs','-R','blocks /fault.bin',str(work/'ioerror.ext4')],stderr=subprocess.DEVNULL).split()
fault_sector=fault_layout['partitions'][0]['start']+int(blocks[0])*8
with VM(work,'writeback-error',system,extra=['-device',controller]) as vm:
    boot(vm)
    vm.qmp('blockdev-add',{'driver':'blkdebug','node-name':'faultdisk',
        'image':{'driver':'raw','file':{'driver':'file','filename':str(fault_disk)}},
        'inject-error':[{'event':'none','iotype':'write','errno':5,'sector':fault_sector,'once':False}]})
    vm.qmp('device_add',{'driver':'usb-storage','id':'bay2','drive':'faultdisk','bus':'xhci.0','port':'2'})
    bay(vm,2,'mounted_read_write')
    capture(vm,"s6-setuidgid fds /bin/bash -c 'printf changed | dd of=/data/fault.bin conv=notrunc status=none'")
    error=rejected(vm,['eject','2'])
    assert 'flush DATA filesystem' in error and 'os error 5' in error,error
    before=bay(vm,2,'error');assert before['mount']=='/data'
    kernel=capture(vm,'dmesg')
    (work/'writeback-kernel.log').write_text(kernel+'\n')
    assert 'I/O error' in kernel or 'error -5' in kernel,kernel
    restart(vm)
    after=bay(vm,2,'error')
    assert after['mount']=='/run/fds/media/02' and after['detail']==before['detail'],(before,after)
    assert 'Shutdown blocked' in rejected(vm,['poweroff'])
    restart(vm)
    assert 'os error 5' in rejected(vm,['poweroff'])
    capture(vm,'test ! -e /run/fds/ejected/02')
    vm.qmp('device_del',{'id':'bay2'});bay(vm,2,'empty')
    vm.send('s6-setuidgid fds fds poweroff');finished(vm)
print('PASS: attached-device writeback EIO is retained across restarts and refuses SAFE and shutdown',flush=True)
link=project/'out/m10-vm-latest'
temporary=link.with_suffix('.next')
temporary.symlink_to(work.name)
temporary.replace(link)
print(f'PASS: M10 ordered shutdown evidence: {work}')
print('SKIP: physical Pi poweroff/reboot, flash-controller durability, battery behavior and sub-second hardware timing targets')