#!/usr/bin/env python3
"""Writable DATA, managed descendants, sync/unmount and hot-removal in ARM VMs."""
import base64
import binascii
import hashlib
import io
import itertools
import json
from pathlib import Path
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time

project=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(project/'tools'))
from image_formats import gpt, LINUX_FILESYSTEM, digest
from vm_test import VM
work=Path(tempfile.mkdtemp(prefix='m7-vm.',dir=project/'out'))
for tool in ['mke2fs','e2fsck','debugfs']:
    assert shutil.which(tool), 'Install the documented host e2fsprogs tools'
with (work/'filesystem-tools.log').open('wb') as log:
    subprocess.run(['mke2fs','-V'],check=True,stdout=log,stderr=subprocess.STDOUT)
controller='qemu-xhci,id=xhci,addr=05.0'
normal=(project/'out/fds-system-cli.img').resolve()

def capture(vm,command):
    # The kernel can write in the middle of a userspace line on ttyAMA0 during
    # deliberate I/O faults. Execute once, then transfer the saved stdout with
    # a checksum. A damaged transfer is retried without repeating the command
    # or discarding any kernel diagnostics from the raw serial log.
    vm.send('('+command+') >/tmp/m7-response; printf "%s" "$?" >/tmp/m7-status; printf "\\nM7_SAVED\\n"')
    vm.expect(rb'^M7_SAVED\r?$')
    for attempt in range(20):
        vm.send('printf "\\nM7_BEGIN\\n"; base64 -w0 /tmp/m7-response; printf "\\n"; sha256sum /tmp/m7-response; cat /tmp/m7-status; printf "\\nM7_END\\n"')
        frame=vm.expect(rb'^M7_BEGIN\r?\n(.*?)\r?\nM7_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',(command,status,payload)
            return payload.decode().strip()
        except (ValueError,binascii.Error):
            continue
    raise AssertionError('Serial transfer repeatedly corrupted; inspect '+str(work/'data.log'))
def query(vm,command='fds --json bays'): return json.loads(capture(vm,command))
def shell(vm,command,marker):
    vm.send(command+f' && printf "\\n{marker}\\n"')
    vm.expect(('^'+marker+r'\r?$').encode())
def wait(vm,command,condition,timeout=40):
    deadline=time.monotonic()+timeout
    while True:
        result=capture(vm,command)
        if condition(result): return result
        if time.monotonic()>deadline: raise AssertionError((command,result))
def bay(vm,n,state):
    value=wait(vm,f'fds --json bay {n}',lambda s:json.loads(s)['bays'][0]['state']==state)
    return json.loads(value)['bays'][0]

with VM(work,'probe',normal,extra=['-device',controller,'-device','usb-kbd,bus=xhci.0,port=3']) as vm:
    vm.expect(rb'FDS> ')
    report=query(vm,'fds --json topology')
    hub=next(d['topology'] for d in report['unmapped'] if '03' in d['interfaces']).rsplit('/',1)[0]
    shell(vm,'test "$(id -u)" = 1000 && touch "$HOME/no-data-check" && test ! -e /data/FDS','M7_EPHEMERAL_HOME')
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))
replacements={'usr/bin/fds-cartridged':(project/'out/fds-cartridged').read_bytes(),'usr/bin/fds':(project/'out/fds').read_bytes(),'etc/fds/bays.toml':config.encode(),'usr/libexec/fds/console-session':b'#!/bin/bash\nexec env HOME=/root bash --login\n'}
with tarfile.open(project/'out/rootfs-cli.tar') as source,tarfile.open(work/'fixture.tar','w',format=tarfile.PAX_FORMAT) as output:
    seen=set()
    for member in source:
        if member.name in replacements:seen.add(member.name);continue
        output.addfile(member,source.extractfile(member) if member.isfile() else None)
    assert seen==replacements.keys()
    for filename,data in {**replacements,'usr/libexec/fds/m7-writer':(project/'out/m7-writer').read_bytes()}.items():
        info=tarfile.TarInfo(filename);info.size=len(data);info.mode=0o755 if filename.startswith('usr/') else 0o644
        output.addfile(info,io.BytesIO(data))
(work/'system').mkdir()
with (work/'image.log').open('wb') as log:
    subprocess.run([str(project/'image/build-system-cartridge'),'--rootfs',str(work/'fixture.tar'),'--output-directory',str(work/'system')],check=True,stdout=log,stderr=subprocess.STDOUT)
system=work/'system/system.img'

def data_image(name):
    root=work/(name+'-files');(root/'FDS').mkdir(parents=True)
    (root/'FDS/CARTRIDGE.TOML').write_text(f'format=1\n[cartridge]\nid="fds.data.{name}"\nname="DATA {name.upper()}"\nclass="data"\nversion="1"\n[media]\nwritable=true\n')
    fs=work/(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/(name+'.img');layout=gpt(disk,[('FDS_DATA',LINUX_FILESYSTEM,fs)])
    return disk,layout
media={name:data_image(name) for name in ['primary','secondary','fault','restart']}
sequence=itertools.count()
with VM(work,'data',system,extra=['-device',controller]) as vm:
    vm.expect(rb'FDS# ')
    def add(name,port):
        node=f'data{next(sequence)}'
        vm.qmp('blockdev-add',{'driver':'raw','node-name':node,'file':{'driver':'file','filename':str(media[name][0])}})
        vm.qmp('device_add',{'driver':'usb-storage','id':f'bay{port}','drive':node,'bus':'xhci.0','port':str(port)})
        return node
    def remove(port):vm.qmp('device_del',{'id':f'bay{port}'});bay(vm,port,'empty')
    add('primary',2);assert bay(vm,2,'mounted_read_write')['mount']=='/data'
    shell(vm,"s6-setuidgid fds /bin/bash -c 'printf persistent > /data/user-file; printf temporary > /home/fds/ephemeral-file'",'M7_USER_WRITES')
    shell(vm,'cd /data; if fds eject 2; then false; else true; fi; cd /','M7_UNMANAGED_BUSY')
    assert bay(vm,2,'mounted_read_write')['mount']=='/data'
    started=query(vm,'s6-setuidgid fds fds --json run 2 -- /usr/libexec/fds/m7-writer')
    assert started['started_pid']>1
    wait(vm,'cat /data/progress 2>/dev/null || printf 0',lambda s:s.isdigit() and int(s)>=128)
    status=capture(vm,'cat /data/process-status')
    assert 'Uid:\t1000\t1000\t1000\t1000' in status and 'CapEff:\t0000000000000000' in status and 'NoNewPrivs:\t1' in status,status
    assert query(vm,'fds --json bay 2')['bays'][0]['consumers']>=2
    shell(vm,'mkdir /run/m7-extra; mount --bind /data /run/m7-extra; if fds eject 2; then false; else true; fi; umount /run/m7-extra','M7_EXTRA_MOUNT_BUSY')
    assert query(vm,'fds --json bay 2')['bays'][0]['consumers']>=2
    shell(vm,'s6-setuidgid fds fds eject 2','M7_SYNCED_SAFE')
    assert bay(vm,2,'safe')['consumers']==0
    shell(vm,'test ! -e /data/FDS && test "$(cat /home/fds/ephemeral-file)" = temporary','M7_HOME_INDEPENDENT')
    shell(vm,'s6-rc -l /run/s6-rc -d change cartridged && s6-rc -l /run/s6-rc -u change cartridged','M7_SAFE_RESTART')
    bay(vm,2,'safe');remove(2)
    # Reinsertion reads persisted data, then eject cleanly for independent fsck.
    add('primary',2);bay(vm,2,'mounted_read_write')
    assert capture(vm,'cat /data/user-file')=='persistent'
    shell(vm,'fds eject 2','M7_REINSERT_SAFE');remove(2)
    print('PASS: sustained writes, UID/capability drop, descendant tracking, TERM escalation, sync/unmount, busy refusal and persistent data',flush=True)

    # Stop the service while two DATA devices enumerate, then inspect one snapshot.
    shell(vm,'s6-rc -l /run/s6-rc -d change cartridged','M7_DAEMON_STOPPED')
    add('primary',2);add('secondary',4)
    wait(vm,"lsblk -nr -o PARTLABEL | grep -c '^FDS_DATA$' || test \"$?\" = 1",lambda s:s=='2')
    shell(vm,'s6-rc -l /run/s6-rc -u change cartridged','M7_DAEMON_STARTED')
    bay(vm,2,'mounted_read_only');bay(vm,4,'mounted_read_only')
    shell(vm,'test ! -e /data/FDS && fds data use 4','M7_EXPLICIT_DATA')
    bay(vm,4,'mounted_read_write')
    shell(vm,'if fds data use 2; then false; else true; fi','M7_NO_DATA_REPLACEMENT')
    shell(vm,'fds eject 2 && fds eject 4','M7_BOTH_SAFE');remove(2);remove(4)
    print('PASS: multiple DATA candidates require selection and cannot replace an active DATA session',flush=True)

    # A crash loses the original syncfs error cursor. A root-owned activation
    # record must prevent the same insertion becoming writable or SAFE again.
    add('restart',2);bay(vm,2,'mounted_read_write')
    shell(vm,"s6-setuidgid fds /bin/bash -c 'printf before-crash > /data/crash-file'",'M10_BEFORE_CRASH')
    shell(vm,'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','M10_CRASH_RESTART')
    quarantined=bay(vm,2,'error')
    assert 'interrupted' in quarantined['detail'] and quarantined['mount']=='/run/fds/media/02',quarantined
    shell(vm,'if fds data use 2; then false; else true; fi; if fds eject 2; then false; else true; fi; test ! -e /run/fds/ejected/02','M10_QUARANTINED')
    assert capture(vm,'cat /run/fds/media/02/crash-file')=='before-crash'
    shell(vm,'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','M10_FAULT_RESTART')
    assert 'interrupted' in bay(vm,2,'error')['detail']
    remove(2)
    print('PASS: interrupted DATA stays read-only and never SAFE across repeated daemon restarts',flush=True)

    # A physical-style pull during dirty writes is a fault, never a successful eject.
    # Preserve kernel I/O diagnostics separately from the serial JSON framing.
    add('fault',2);bay(vm,2,'mounted_read_write')
    query(vm,'fds --json run 2 -- /usr/libexec/fds/m7-writer')
    wait(vm,'cat /data/progress 2>/dev/null || printf 0',lambda s:s.isdigit() and int(s)>=16)
    remove(2)
    shell(vm,'if fds eject 2; then false; else true; fi','M7_REMOVED_NOT_SAFE')
    shell(vm,"grep -q 'populated 0' /sys/fs/cgroup/fds/bay02/cgroup.events && test ! -e /data/FDS",'M7_REMOVAL_CLEANUP')
    diagnostics=capture(vm,'dmesg')
    (work/'surprise-removal-kernel.log').write_text(diagnostics+'\n')
    assert 'error -5' in diagnostics or 'I/O error' in diagnostics, 'Fault test did not observe a kernel storage error'
    print('PASS: surprise removal stops managed descendants and does not report SAFE',flush=True)

# Verify unmounted ext4 using independent filesystem tools, without a host mount.
for name in ['primary','secondary']:
    disk,layout=media[name];part=layout['partitions'][0];output=work/(name+'-after.ext4')
    with disk.open('rb') as source,output.open('wb') as target:
        source.seek(part['start']*512);remaining=part['payload_bytes']
        while remaining:
            data=source.read(min(1024*1024,remaining));assert data;target.write(data);remaining-=len(data)
    with (work/(name+'-fsck.log')).open('wb') as log:
        subprocess.run(['e2fsck','-fn',str(output)],check=True,stdout=log,stderr=subprocess.STDOUT)
    if name=='primary':
        extracted=work/'stress-after.bin'
        subprocess.run(['debugfs','-R',f'dump /stress.bin {extracted}',str(output)],check=True,stdout=subprocess.DEVNULL)
        expected=hashlib.sha256(bytes(range(256))*(64*1024*1024//256)).hexdigest()
        assert extracted.stat().st_size==64*1024*1024 and digest(extracted)==expected
print('PASS: both safely ejected DATA filesystems pass e2fsck; 64 MiB sustained-write payload matches exactly')
link=project/'out/m7-vm-latest';temporary=link.with_suffix('.next');temporary.symlink_to(work.name);temporary.replace(link)
print(f'PASS: M7 writable DATA integration: {work}')
print('SKIP: physical storage power-loss behavior and Pi eject/shutdown latency')