"""Actual USB boot, ordinary FAT program loading and two cold boots of the same writable medium."""
import argparse
import hashlib
import json
from pathlib import Path
import re
import shutil
import socket
import subprocess
import time
from PIL import Image
from TestBoot import ROOT, Qmp, digest, read_json
from IndependentBingEvidence import json_blocks
from IndependentConsoleEvidence import terminal_command_output, FAILURE_MARKERS


def verify_storage(image, edition=None, force_other_disks=False):
    image=Path(image).resolve();manifest=read_json(str(image)+'.json');edition=edition or manifest['Edition']
    assert manifest.get('UsbWorkspace') is True and edition==manifest['Edition'] and edition in ('Console','Desktop')
    proof=ROOT/'build/IndependentKernel'/('UsbStorage-'+str(time.time_ns()));proof.mkdir()
    disk=proof/'WritableUsb.img';shutil.copyfile(image,disk)
    result={'Kind':'XfeNativeUsbPersistentWorkspaceProof','Passed':False,'Edition':edition,
        'SourceImagePath':str(image),'SourceImageSha256':digest(image),'SourceManifestSha256':digest(str(image)+'.json'),
        'KernelSha256':manifest['KernelSha256'],'Input':'Actual QMP USB HID packets only',
        'HostCommandExecution':False,'HostWorkspaceInjection':False,'PhysicalHardwareAcceptancePassed':False,
        'ImageWrites':'Private copy of the exact source image; original image remains unchanged',
        'Boots':[],'EvidenceFiles':[],
        'SourceHashes':{str(p.relative_to(ROOT)).replace('\\','/'):digest(p) for p in [Path(__file__),ROOT/'Tools/CSharpBoot/ReadImage.py']}}
    result['SourceHashes'].update(manifest['IndependentKernelBuild']['SourceSha256'])
    result['SourceHashes'].update(manifest['StorageSourceSha256'])
    other_disks=force_other_disks or manifest['IndependentKernelBuild']['Configuration'].get('OtherDiskAccessEnabled') is True
    fixture_manifest=None;other_sources=[]
    if other_disks:
        generator=ROOT/'Tools/StorageFixtures/bin/Release/net10.0/StorageFixtures.dll'
        assert generator.is_file(),'Build the locked StorageFixtures tool before actual disk acceptance'
        seed=proof/'OtherDiskSeeds'
        with (proof/'FixtureGeneration.log').open('w',encoding='utf-8') as log:
            completed=subprocess.run(['dotnet',str(generator),'create',str(seed),str(ROOT/'Samples/XFESystem.UsbHello/bin/Release/net10.0/XFESystem.UsbHello.dll')],cwd=ROOT,stdout=log,stderr=subprocess.STDOUT,timeout=90,creationflags=getattr(subprocess,'CREATE_NO_WINDOW',0))
        assert completed.returncode==0,'Independent filesystem fixture generation failed'
        fixture_manifest=read_json(seed/'Manifest.json');assert fixture_manifest['Passed'] and fixture_manifest['PhysicalDiskAccess'] is False
        for item in fixture_manifest['Images']:
            source=Path(item['Path']);assert digest(source)==item['Sha256']
            target=proof/source.name;shutil.copyfile(source,target);other_sources.append((source,target,item))
        result.update(OtherDiskAcceptanceRequired=True,OtherDiskAcceptancePassed=False,FixtureManifestSha256=digest(seed/'Manifest.json'))
        for name in ('StorageFixtures.csproj','FixtureDisk.cs','Program.cs','StorageReadback.cs','packages.lock.json'):
            path=ROOT/'Tools/StorageFixtures'/name;result['SourceHashes'][path.relative_to(ROOT).as_posix()]=digest(path)
        result['FixtureGeneratorSha256']=digest(generator)
        result['FixtureGeneratorPath']=str(generator)
    token='UsbSaved'+str(time.time_ns())[-12:]
    def ref(path):return {'Path':path.relative_to(proof).as_posix(),'Sha256':digest(path)}
    try:
        for boot_number in (1,2):
            folder=proof/('Boot'+str(boot_number));folder.mkdir();shutil.copyfile(ROOT/'.tools/firmware/edk2-i386-vars.fd',folder/'Vars.fd')
            with socket.socket() as reservation:reservation.bind(('127.0.0.1',0));port=reservation.getsockname()[1]
            command=[str(ROOT/'.tools/qemu/qemu-system-x86_64.exe'),'-machine','q35,i8042=off','-accel','tcg','-cpu','max','-smp','8','-m','1024M','-vga','std',
                '-global','VGA.vgamem_mb=64','-global','VGA.edid=on','-global','VGA.xres=1280','-global','VGA.yres=800',
                '-global','VGA.xmax=1280','-global','VGA.ymax=800','-display','none','-net','none','-rtc','base=utc',
                '-drive',f'if=pflash,format=raw,readonly=on,file={ROOT/".tools/firmware/edk2-x86_64-code.fd"}',
                '-drive',f'if=pflash,format=raw,file={folder/"Vars.fd"}',
                '-device','qemu-xhci,id=Usb','-drive',f'if=none,id=BootUsb,format=raw,file={disk}',
                '-device','usb-storage,drive=BootUsb,bus=Usb.0,bootindex=1',
                '-device','usb-kbd,bus=Usb.0','-device','usb-mouse,bus=Usb.0',
                '-serial',f'file:{folder/"Serial.log"}','-qmp',f'tcp:127.0.0.1:{port},server=on,wait=off','-no-reboot','-no-shutdown']
            if other_disks:
                command+=['-drive',f'if=none,id=OtherNvme,format=raw,file={proof/"NvmeNtfs.img"}','-device','nvme,drive=OtherNvme,serial=XFE_OTHER_NVME',
                    '-drive',f'if=none,id=OtherSata,format=raw,file={proof/"SataExtended.img"}','-device','ide-hd,drive=OtherSata,bus=ide.0',
                    '-drive',f'if=none,id=OtherUsb,format=raw,file={proof/"UsbExFat.img"}','-device','usb-storage,drive=OtherUsb,bus=Usb.0',
                    '-drive',f'if=none,id=Other4K,format=raw,file={proof/"Nvme4K.img"}','-device','nvme,drive=Other4K,serial=XFE_OTHER_4K,logical_block_size=4096,physical_block_size=4096']
            boot={'Number':boot_number,'QemuArguments':command,'Commands':[]};result['Boots'].append(boot)
            qemu=qmp=None;selected=[]
            def text():
                path=folder/'Serial.log';value=path.read_text(encoding='utf-8',errors='replace') if path.exists() else ''
                assert qemu.poll() is None,'USB guest exited unexpectedly'
                assert not any(marker in value for marker in FAILURE_MARKERS),value[-3000:]
                return value
            def wait(marker,after=0,timeout=50):
                end=time.monotonic()+timeout
                while time.monotonic()<end:
                    if marker in text()[after:]:return
                    time.sleep(.1)
                raise AssertionError('USB workspace timeout: '+marker+'\n'+text()[-2000:])
            def key(chord):
                qmp.call('human-monitor-command',{'command-line':'sendkey '+chord+(' 20' if other_disks else '')});time.sleep(.045 if other_disks else .11)
            def type_text(value):
                special={':':'shift-semicolon','.':'dot','/':'slash','-':'minus','_':'shift-minus',' ':'spc'}
                for char in value:key(special.get(char,'shift-'+char.lower() if char.isupper() else char))
            def execute(value,required=()):
                current=text();previous=json_blocks(current,'XFE_CONSOLE_SNAPSHOT' if edition=='Console' else 'XFE_TERMINAL_SNAPSHOT')
                echo_prompt=previous[-1]['Value']['Prompt'] if previous else None
                after=len(current);type_text(value);key('ret');deadline=time.monotonic()+35;last=''
                while time.monotonic()<deadline:
                    packets=json_blocks(text(),'XFE_CONSOLE_SNAPSHOT' if edition=='Console' else 'XFE_TERMINAL_SNAPSHOT',after)
                    if packets:
                        packet=packets[-1]
                        try:
                            output=terminal_command_output(packet['Value'],value,echo_prompt)
                            assert all(line in output for line in required),repr(output)
                            selected.append(packet);boot['Commands'].append({'Command':value,'PostCommandOutput':output,'Passed':True});return
                        except AssertionError as error:last=str(error)
                    time.sleep(.1)
                raise AssertionError('USB command did not render: '+value+': '+last)
            def screenshot(name):
                ppm=folder/(name+'.ppm');png=folder/(name+'.png');qmp.call('screendump',{'filename':str(ppm)})
                with Image.open(ppm) as frame:assert frame.size==(1280,800);frame.convert('RGB').save(png)
                result['EvidenceFiles'].append(ref(png));return ref(png)
            def terminal():
                if edition=='Desktop':
                    after=len(text());key('ctrl-esc');key('down');key('down');key('down');key('ret');wait('XFE_DESKTOP_ACTIVE|Terminal',after)
            try:
                with (folder/'Qemu.log').open('w') as log:qemu=subprocess.Popen(command,stdout=log,stderr=subprocess.STDOUT,creationflags=getattr(subprocess,'CREATE_NO_WINDOW',0))
                deadline=time.monotonic()+15
                while qmp is None:
                    assert qemu.poll() is None,(folder/'Qemu.log').read_text()
                    try:
                        qmp=Qmp(port)
                        qmp.connection.setsockopt(socket.IPPROTO_TCP,socket.TCP_NODELAY,1)
                    except OSError:assert time.monotonic()<deadline;time.sleep(.1)
                wait('XFE_INDEPENDENT_'+('CONSOLE' if edition=='Console' else 'DESKTOP')+'_FIRST_FRAME',timeout=120)
                log=text();assert 'XFE_WORKSPACE|Mode|USB|Persistent|1|' in log and 'XFE_USB_STORAGE_READY|Matches|1|' in log
                match=re.search(r'XFE_WORKSPACE\|Mode\|USB\|Persistent\|1\|Generation\|(\d+)',log);assert match
                boot['LoadedGeneration']=int(match[1]);boot['BootScreenshot']=screenshot('Boot')
                if boot_number==2:
                    assert boot['LoadedGeneration']>0,'Cold boot did not load a committed physical workspace'
                    if edition=='Desktop':
                        # UI diagnostic snapshots are deliberately asynchronous.
                        # Wait for the actual restored state, then read fresh data;
                        # the earlier first-frame log cannot establish appearance.
                        wait('XFE_DESKTOP_APPEARANCE|Dark|')
                        restored=re.findall(r'^XFE_DESKTOP_APPEARANCE\|Dark\|[0-9A-F]{6}\|Persistent\|1\r?$',text(),re.M)
                        assert restored,'Cold boot did not report the restored persistent dark appearance'
                if boot_number==1 and edition=='Desktop':
                    after=len(text());key('ctrl-esc')
                    for _ in range(4):key('down')
                    key('ret');wait('XFE_DESKTOP_ACTIVE|Settings',after);key('right');wait('XFE_DESKTOP_APPEARANCE|Dark|',after)
                    boot['SettingsScreenshot']=screenshot('Settings-Saved')
                terminal();execute('help',('Storage: persistent workspace',))
                if boot_number==1:
                    execute('mf UsbKeep.txt');execute('write UsbKeep.txt '+token)
                    execute('mv UsbKeep.txt UsbRenamed.txt');execute('cat UsbRenamed.txt',(token,))
                    execute('write /System/Configuration/UsbSetting.txt '+token)
                    execute('run /UsbFiles/UsbHello.dll',('XFE_USB_EXTERNAL_APP_OK','Managed ConsoleIL exit: 0'))
                    execute('cp /UsbFiles/UsbHello.dll UsbCopied.dll')
                else:
                    execute('cat UsbRenamed.txt',(token,));execute('cat /System/Configuration/UsbSetting.txt',(token,))
                    execute('run UsbCopied.dll',('XFE_USB_EXTERNAL_APP_OK','Managed ConsoleIL exit: 0'))
                    execute('run /UsbFiles/UsbHello.dll',('XFE_USB_EXTERNAL_APP_OK','Managed ConsoleIL exit: 0'))
                if other_disks:
                    mounts=[{'Path':m[1],'FileSystem':m[2],'Writable':m[3]=='1','Label':m[4]} for m in re.finditer(r'^XFE_VOLUME_MOUNT\|([^|\r\n]+)\|FS\|(\w+)\|Writable\|([01])\|Label\|([^\r\n]*)',text(),re.M)]
                    assert len(mounts)==5 and all(m['Writable'] for m in mounts),mounts
                    assert sorted(m['FileSystem'] for m in mounts)==['FAT','NTFS','NTFS','NTFS','exFAT'],mounts
                    boot['OtherVolumes']=mounts;execute('ls /Disks')
                    if boot_number==1:(proof/'ExpectedOtherVolumes.json').write_text(json.dumps(mounts,indent=2),encoding='utf-8')
                    for mount in mounts:
                        root=mount['Path'].replace('X:\\','/').replace('\\','/');saved=token+'_'+mount['Path'].rsplit('\\',1)[-1]
                        execute('cd '+root)
                        if boot_number==1:
                            execute('cat Readme.txt')
                            execute('md XfeProof')
                            execute('write XfeProof/Keep.txt '+saved)
                            execute('append XfeProof/Keep.txt _appended')
                            execute('mv XfeProof/Keep.txt XfeProof/Renamed.txt')
                            execute('write EditMe.txt '+saved)
                            execute('del DeleteMe.txt')
                            execute('cp /UsbFiles/UsbHello.dll XfeProof/External.dll')
                            execute('cp Pattern.bin XfeProof/PatternCopy.bin')
                        execute('cat EditMe.txt',(saved,))
                        execute('cat XfeProof/Renamed.txt',(saved+'_appended',))
                        execute('run XfeProof/External.dll',('XFE_USB_EXTERNAL_APP_OK','Managed ConsoleIL exit: 0'))
                    boot['OtherDiskReadWritePassed']=True
                    if edition=='Desktop':
                        # Opening a folder activates Files before the terminal's
                        # next render; require that actual UI transition instead.
                        after=len(text());type_text('open /Disks');key('ret');wait('XFE_DESKTOP_ACTIVE|Files',after)
                        boot['DiskListScreenshot']=screenshot('Graphical-Disk-List')
                boot['FinalScreenshot']=screenshot('Workspace-Verified');boot['Passed']=True
                raw=folder/'RenderedCommands.log';raw.write_text('\n'.join(p['Raw'] for p in selected),encoding='utf-8');result['EvidenceFiles'].append(ref(raw))
                qmp.call('quit');qemu.wait(timeout=10);boot['QemuExitCode']=qemu.returncode;assert qemu.returncode==0
            finally:
                if qmp is not None:qmp.connection.close()
                if qemu is not None and qemu.poll() is None:qemu.terminate();qemu.wait(timeout=10)
                if (folder/'Serial.log').exists():result['EvidenceFiles'].append(ref(folder/'Serial.log'))
        original=image.read_bytes();actual=disk.read_bytes();offset=manifest['WorkspaceFirstLba']*512;size=manifest['WorkspaceBytes']
        assert len(original)==len(actual) and original[:offset+512]==actual[:offset+512] and original[offset+size:]==actual[offset+size:]
        assert original[offset+512:offset+size]!=actual[offset+512:offset+size],'No data changed on the actual USB disk'
        assert all(digest(ROOT/p)==h for p,h in result['SourceHashes'].items())
        if other_disks:
            result['OtherDiskImages']=[]
            for source,target,item in other_sources:
                assert digest(source)==item['Sha256'] and digest(target)!=item['Sha256']
                result['OtherDiskImages'].append({'Kind':item['Kind'],'SourceSha256':item['Sha256'],'ActualSha256':digest(target),'Path':str(target)})
            verification=proof/'OtherDiskReadback.json'
            completed=subprocess.run(['dotnet',str(generator),'verify',str(proof),token],cwd=ROOT,capture_output=True,timeout=90,creationflags=getattr(subprocess,'CREATE_NO_WINDOW',0))
            (proof/'ReadbackTool.stdout.log').write_bytes(completed.stdout);(proof/'ReadbackTool.stderr.log').write_bytes(completed.stderr)
            assert completed.returncode==0,(completed.stderr or completed.stdout).decode('utf-8',errors='replace')
            check=read_json(verification);assert check['Passed'] and check['ColdBootFilesVerified'] and check['OnlyPartitionSectorsChanged'] and check['AllVolumesClean']
            result['EvidenceFiles'].extend([ref(verification),ref(proof/'FixtureGeneration.log'),ref(seed/'Manifest.json'),
                ref(proof/'ExpectedOtherVolumes.json'),ref(proof/'ReadbackTool.stdout.log'),ref(proof/'ReadbackTool.stderr.log')])
            assert digest(generator)==result['FixtureGeneratorSha256']
            result.update(OtherDiskAcceptancePassed=True,NvmeAhciOtherUsbReadWritePassed=True,Nvme4096ReadModifyWritePassed=True,
                NtfsFatExfatReadWritePassed=True,OtherDisksPersistedAcrossColdBoot=True,OtherDiskProgramExecutionPassed=True,
                OtherDiskHostReadback={'Path':verification.name,'Sha256':digest(verification)})
        result.update(Passed=True,SourceImageUnchanged=digest(image)==result['SourceImageSha256'],OnlyWorkspaceSectorsChanged=True,
            FilePersistedAcrossColdBoot=True,ConfigurationPersistedAcrossColdBoot=True,AppearancePersistedAcrossColdBoot=edition=='Desktop',
            StandardNetPeReadFromFatAndExecuted=True,ImportedBinaryCopiedToWorkspaceAndExecutedAfterReboot=True,
            WritableImageSha256=digest(disk),WritableImagePath=str(disk),SourcesUnchanged=True,ProofDirectory=str(proof))
    except Exception as error:
        result['Failure']=str(error)
    finally:
        (proof/'Result.json').write_text(json.dumps(result,ensure_ascii=False,indent=2),encoding='utf-8')
    return result,proof/'Result.json'


if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('image',type=Path);parser.add_argument('--other-disks',action='store_true');args=parser.parse_args()
    result,path=verify_storage(args.image,force_other_disks=args.other_disks);print(json.dumps({'Passed':result['Passed'],'Receipt':str(path),'Failure':result.get('Failure')},ensure_ascii=True))
    raise SystemExit(0 if result['Passed'] else 1)