#!/usr/bin/env python3
"""Actual ARM WindowMaker input, ENVIRONMENT lifecycle, PROGRAM launch and DHCP."""
import base64
import io
import gzip
import itertools
import json
from pathlib import Path
import re
import struct
import subprocess
import sys
import tarfile
import tempfile
import time
import zlib

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='m8-vm.', dir=project/'out'))
private_network=subprocess.check_output([str(project/'tools/in-void'),'--isolated-network','readlink','/proc/self/ns/net'],text=True).strip()
assert private_network != str(Path('/proc/self/ns/net').readlink()), 'VM transport must not use the host network namespace'
(work/'network-isolation.txt').write_text('QEMU runs in a private network namespace; virtual DHCP/DNS has no host or Internet route.\n')
controller = 'qemu-xhci,id=xhci,addr=05.0'

def capture(vm, command):
    vm.send('printf "\\nM8_BEGIN\\n"; '+command+'; printf "\\nM8_END\\n"')
    return vm.expect(rb'^M8_BEGIN\r?\n(.*?)\r?\nM8_END\r?$').group(1).decode().strip()
def query(vm, command='fds --json profiles'): 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=45):
    deadline=time.monotonic()+timeout
    while True:
        result=capture(vm, command)
        if condition(result): return result
        if time.monotonic()>deadline:
            details=capture(vm,'cat /run/log/network/current 2>/dev/null; ip -brief address')
            raise AssertionError((command,result,details))
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 desktop(vm, enabled):
    if enabled:
        result=capture(vm, 's6-svwait -U -t 25000 /run/service/desktop-session; printf "WAIT_STATUS:%s" "$?"')
        if result != 'WAIT_STATUS:0':
            raise AssertionError((result,capture(vm,'cat /run/log/xserver/current /run/log/desktop/current /run/log/cartridged/current')))
    def condition(s):
        p=json.loads(s)['profiles']
        if p['error']: raise AssertionError(p)
        return p['desktop']==('windowmaker' if enabled else 'cli') and (p['ready_ns'] is not None if enabled else True)
    return json.loads(wait(vm,'fds --json profiles',condition))['profiles']
def image(name, replacements):
    replacements={**replacements, **{f'usr/bin/{binary}':(project/'out'/binary).read_bytes() for binary in ['fds','fds-program','fds-control','fds-cartridged','fds-profile']}}
    with tarfile.open(project/'out/rootfs-development.tar') as source, tarfile.open(work/(name+'.tar'),'w',format=tarfile.PAX_FORMAT) as output:
        assert source.extractfile('usr/share/fds/image-profile').read().strip()==b'development', 'Build PROFILE=development before this test'
        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(),(seen,replacements.keys())
        for filename,data in replacements.items():
            info=tarfile.TarInfo(filename);info.size=len(data);info.mode=0o755 if filename.startswith('usr/') else 0o644
            output.addfile(info,io.BytesIO(data))
    destination=work/name;destination.mkdir()
    with (work/(name+'-image.log')).open('wb') as log:
        subprocess.run([str(project/'image/build-system-cartridge'),'--rootfs',str(work/(name+'.tar')),'--output-directory',str(destination)],check=True,stdout=log,stderr=subprocess.STDOUT)
    return destination/'system.img'
root_console=b'#!/bin/bash\n# Isolated test image only.\nexec env HOME=/root bash --login\n'
base={'usr/libexec/fds/console-session':root_console,'etc/fds/xserver':b'xvfb\n'}
probe=image('probe-system',base)
with VM(work,'probe',probe,extra=['-device',controller,'-device','usb-kbd,bus=xhci.0,port=3']) as vm:
    vm.expect(rb'FDS# ')
    shell(vm,'fds-boottrace mark console-ready','M8_TEST_CONSOLE_OBSERVED')
    topology=query(vm,'fds --json topology')
    hub=next(d['topology'] for d in topology['unmapped'] if '03' in d['interfaces']).rsplit('/',1)[0]
    shell(vm,'ip -o link show dev lo | grep -q UP','M8_LOCAL_LOOPBACK')
    p=query(vm)['profiles'];assert p['desktop']=='cli' and p['network']==[] and p['ready_ns'] is None
    shell(vm,'! pgrep -x Xvfb && ! pgrep -x wmaker && ! pgrep -x dhcpcd','M8_OPTIONAL_SERVICES_OFF')
    shell(vm,'s6-setuidgid fds fds profile activate windowmaker','M8_ACTIVATION_REQUESTED')
    try:
        p=desktop(vm,True)
    except BaseException:
        print(capture(vm,'cat /run/log/xserver/current /run/log/desktop/current /run/log/cartridged/current'),flush=True)
        raise
    shell(vm,'test ! -d /home/fds/.cache/fontconfig','M8_PREBUILT_FONT_CACHE')
    shell(vm,'cmp /etc/WindowMaker/WindowMaker /usr/share/fds/eink/WindowMaker && cmp /etc/WindowMaker/WMRootMenu /usr/share/fds/eink/WMRootMenu && grep -q "FDS Control" /etc/WindowMaker/WMRootMenu','FDS_GLOBAL_RETRO_DEFAULTS')
    (work/'activation.json').write_text(json.dumps(p,indent=2)+'\n')
    assert p['ready_ns']>=p['activation_ns']>0
    shell(vm,'export DISPLAY=:0 XAUTHORITY=/run/fds/x11/authority; s6-setuidgid fds xdpyinfo >/tmp/display-info; ! grep -q "COMPOSITE" /tmp/display-info','M8_X11_AUTHENTICATED')
    shell(vm,'if env XAUTHORITY=/dev/null xdpyinfo >/dev/null 2>&1; then false; else true; fi','M8_X11_REJECTS_NO_COOKIE')
    panel=wait(vm,'s6-setuidgid fds xdotool search --onlyvisible --name "^FDS Control$"',lambda s:s.isdigit())
    shell(vm,'test "$(ps -o uid= -C fds-control | xargs)" = 1000','FDS_CONTROL_UNPRIVILEGED')
    window=wait(vm,'s6-setuidgid fds xdotool search --onlyvisible --name "^FDS Terminal$"',lambda s:s.isdigit())
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {window} type --clearmodifiers "printf M8_INPUT_WORKED > /home/fds/desktop-input"; s6-setuidgid fds xdotool key --clearmodifiers Return','M8_X11_INPUT_SENT')
    wait(vm,'cat /home/fds/desktop-input 2>/dev/null',lambda s:s=='M8_INPUT_WORKED')
    shell(vm,'s6-setuidgid fds xwd -root -silent -out /home/fds/desktop.xwd','M8_SCREENSHOT_SAVED')
    encoded=capture(vm,'gzip -c /home/fds/desktop.xwd | base64 -w0')
    (work/'desktop.xwd').write_bytes(gzip.decompress(base64.b64decode(encoded)))
    shell(vm,'s6-setuidgid fds fds profile deactivate && ! pgrep -x Xvfb && ! pgrep -x wmaker && ! pgrep -x xterm && pgrep -x dasungd && test "$(cat /home/fds/desktop-input)" = M8_INPUT_WORKED','M8_BACK_TO_CONSOLE')
    desktop(vm,False)
    shell(vm,'s6-setuidgid fds fds profile activate windowmaker','M8_REACTIVATE')
    second=desktop(vm,True);assert second['activation_ns']>p['activation_ns']
    report=query(vm,'fds --json boot-profile');assert report['events_ns']['desktop-ready']==p['ready_ns']
    shell(vm,'fds profile deactivate','M8_STOP_SECOND_SESSION')
    print('PASS: optional desktop, authenticated X11, WindowMaker input, terminal, clean stop/restart and measured activation',flush=True)

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))
fixture=image('mapped-system',{**base,'etc/fds/bays.toml':config.encode()})
def cartridge(name, manifest, label, files=None):
    with tarfile.open(work/(name+'.tar'),'w') as output:
        for folder in ['FDS','app','app/bin','app/lib','app/share']:
            info=tarfile.TarInfo(folder);info.type=tarfile.DIRTYPE;info.mode=0o755;output.addfile(info)
        for filename,data in {'FDS/CARTRIDGE.TOML':manifest,**(files or {})}.items():
            info=tarfile.TarInfo(filename);info.mode=0o755 if filename.startswith('app/bin/') else 0o644
            if isinstance(data,tuple):
                info.type=tarfile.SYMTYPE;info.linkname=data[0];output.addfile(info)
            else:
                info.size=len(data);output.addfile(info,io.BytesIO(data))
    erofs=work/(name+'.erofs')
    subprocess.run([str(project/'tools/in-image-tools'),'mkfs.erofs','-T','0','--tar=f',str(erofs),str(work/(name+'.tar'))],check=True,stdout=subprocess.DEVNULL)
    disk=work/(name+'.img');gpt(disk,[(label,LINUX_FILESYSTEM,erofs)]);return disk
manifest=(project/'tests/fixtures/manifests/windowmaker.toml').read_bytes()
env=cartridge('environment',manifest,'FDS_ENVIRONMENT')
unknown=cartridge('unknown-profile',manifest.replace(b'profile = "windowmaker"',b'profile = "unknown"'),'FDS_ENVIRONMENT')
program_manifest=b'format=1\n[cartridge]\nid="fds.program.test"\nname="PROGRAM TEST"\nclass="program"\nversion="1"\n[media]\nwritable=false\n'
program=cartridge('program',program_manifest,'FDS_PROGRAM',{'app/bin/escape':('/usr/bin/id',),'app/bin/check':b'#!/bin/bash\nset -eu\nid > "$HOME/program-identity"\nprintf "%s\\n" "$FDS_APP" "$LD_LIBRARY_PATH" > "$HOME/program-paths"\nexec /usr/bin/tail -f /dev/null\n'})
sequence=itertools.count()
with VM(work,'cartridges',fixture,extra=['-device',controller,'-netdev','user,id=net']) as vm:
    vm.expect(rb'FDS# ')
    shell(vm,'fds-boottrace mark console-ready','M8_TEST_CONSOLE_OBSERVED')
    def add(path,port):
        node=f'media{next(sequence)}'
        vm.qmp('blockdev-add',{'driver':'raw','node-name':node,'read-only':True,'file':{'driver':'file','filename':str(path)}})
        vm.qmp('device_add',{'driver':'usb-storage','id':f'bay{port}','drive':node,'bus':'xhci.0','port':str(port)})
    def remove(port):vm.qmp('device_del',{'id':f'bay{port}'});bay(vm,port,'empty')
    add(unknown,2);bay(vm,2,'mounted_read_only');assert query(vm)['profiles']['desktop']=='cli';remove(2)
    add(env,2);bay(vm,2,'mounted_read_only');assert desktop(vm,True)['environment_bay']==2
    shell(vm,'s6-setuidgid fds fds eject 2 && ! pgrep -x wmaker && pgrep -x dasungd','M8_ENV_SAFE_RETURN')
    bay(vm,2,'safe');desktop(vm,False);remove(2)
    add(env,2);bay(vm,2,'mounted_read_only');desktop(vm,True)
    shell(vm,'fds profile deactivate; fds rescan >/dev/null','M8_ENV_MANUAL_DEACTIVATION')
    desktop(vm,False);remove(2)
    add(env,2);bay(vm,2,'mounted_read_only');desktop(vm,True);remove(2);desktop(vm,False)
    add(program,4);entry=bay(vm,4,'mounted_read_only');assert entry['mount']=='/run/fds/apps/fds.program.test'
    shell(vm,'test ! -e /home/fds/program-identity && ! /run/fds/apps/fds.program.test/app/bin/check','M8_NO_AUTORUN')
    shell(vm,'fds profile activate windowmaker','FDS_CONTROL_START_DESKTOP')
    desktop(vm,True)
    shell(vm,'export DISPLAY=:0 XAUTHORITY=/run/fds/x11/authority','FDS_CONTROL_DISPLAY')
    panel=wait(vm,'s6-setuidgid fds xdotool search --onlyvisible --name "^FDS Control$"',lambda s:s.isdigit())
    wait(vm,'s6-setuidgid fds xdotool search --onlyvisible --name "^FDS Terminal$"',lambda s:s.isdigit())
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel} windowraise {panel} mousemove --window {panel} 70 218 click 1','FDS_CONTROL_SELECT_BAY')
    def panel_state():
        output=capture(vm,f's6-setuidgid fds xprop -id {panel} _FDS_CONTROL_STATE')
        return json.loads(json.loads(output.split(' = ',1)[1]))
    deadline=time.monotonic()+30
    while True:
        snapshot=panel_state()
        if snapshot['bay']==4 and snapshot['commands']>0 and not snapshot['busy']: break
        assert time.monotonic()<deadline,snapshot
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel} windowraise {panel} mousemove --window {panel} 420 518 click 1','FDS_CONTROL_RUN_CLICK')
    wait(vm,'cat /home/fds/program-identity 2>/dev/null',lambda s:'uid=1000(fds)' in s)
    assert query(vm,'fds --json bay 4')['bays'][0]['consumers']>0
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel}','FDS_CONTROL_SHOW_PANEL')
    shell(vm,'s6-setuidgid fds xwd -root -silent -out /home/fds/control.xwd','FDS_CONTROL_SCREENSHOT')
    (work/'legacy-control.xwd').write_bytes(gzip.decompress(base64.b64decode(capture(vm,'gzip -c /home/fds/control.xwd | base64 -w0'))))
    wait(vm,'cat /home/fds/program-identity 2>/dev/null',lambda s:'uid=1000(fds)' in s)
    assert '/run/fds/apps/fds.program.test/app/lib' in capture(vm,'cat /home/fds/program-paths')
    shell(vm,'if fds run 4 -- escape; then false; else true; fi','M8_PROGRAM_ESCAPE_REJECTED')
    shell(vm,'if fds run 4 -- ../check; then false; else true; fi','M8_PROGRAM_PATH_REJECTED')
    deadline=time.monotonic()+30
    while panel_state()['busy']:
        assert time.monotonic()<deadline
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel} windowraise {panel} mousemove --window {panel} 612 518 click 1','FDS_CONTROL_EJECT_CLICK')
    assert bay(vm,4,'safe')['consumers']==0;remove(4)
    # Launch a format-2, Void-built command through the real panel as well.
    built=Path((project/'out/workstation-images-current.txt').read_text().strip())
    assert json.loads((built/'acceptance.json').read_text())['status']=='passed'
    add(built/'software.img',4);bay(vm,4,'mounted_read_only')
    deadline=time.monotonic()+30
    while True:
        snapshot=panel_state()
        if snapshot['commands']>=2 and not snapshot['busy']: break
        assert time.monotonic()<deadline,snapshot
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel} windowraise {panel} mousemove --window {panel} 420 518 click 1','FDS_CONTROL_RUN_VOID')
    hello_window=wait(vm,'s6-setuidgid fds xdotool search --onlyvisible --name "^demo.hello:hello$"',lambda s:s.isdigit())
    shell(vm,f's6-setuidgid fds xdotool windowsize {hello_window} 650 360 windowmove {hello_window} 930 70 windowactivate --sync {panel}','FDS_CONTROL_VOID_RESULT')
    shell(vm,'s6-setuidgid fds xwd -root -silent -out /home/fds/control.xwd','FDS_CONTROL_VOID_SCREENSHOT')
    (work/'control.xwd').write_bytes(gzip.decompress(base64.b64decode(capture(vm,'gzip -c /home/fds/control.xwd | base64 -w0'))))
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel} windowraise {panel} key r','FDS_CONTROL_KEYBOARD_RESCAN')
    deadline=time.monotonic()+30
    while True:
        snapshot=panel_state()
        if not snapshot['busy'] and snapshot['status']=='Cartridge inventory refreshed.': break
        assert time.monotonic()<deadline,snapshot
    shell(vm,f's6-setuidgid fds xdotool windowactivate --sync {panel} windowraise {panel} mousemove --window {panel} 612 518 click 1','FDS_CONTROL_EJECT_VOID')
    assert bay(vm,4,'safe')['consumers']==0;remove(4)
    shell(vm,'fds profile deactivate && ! pgrep -x fds-control','FDS_CONTROL_SESSION_STOPPED')
    shell(vm,'modprobe cdc_ether','M8_NETWORK_DRIVER')
    vm.qmp('device_add',{'driver':'usb-net','id':'ethernet','netdev':'net','bus':'xhci.0','port':'3'})
    wait(vm,'fds --json profiles',lambda s:len(json.loads(s)['profiles']['network'])==1)
    address=json.loads(wait(vm,'ip -j -4 address show',lambda s:any(a.get('local')=='10.0.2.15' for i in json.loads(s) for a in i['addr_info'])))
    (work/'dhcp.json').write_text(json.dumps(address,indent=2)+'\n')
    wait(vm,'cat /run/fds/resolv.conf',lambda s:'nameserver 10.0.2.3' in s)
    shell(vm,'s6-setuidgid fds fds network off && ! pgrep -x dhcpcd','M8_NETWORK_OFF')
    assert not any('UP' in i['flags'] for i in json.loads(capture(vm,'ip -j link show')) if i['ifname']!='lo')
    shell(vm,'s6-setuidgid fds fds network on','M8_NETWORK_EXPLICIT')
    wait(vm,'ip -j -4 address show',lambda s:any(a.get('local')=='10.0.2.15' for i in json.loads(s) for a in i['addr_info']))
    shell(vm,'fds network off','M8_NETWORK_END')
    vm.qmp('device_del',{'id':'ethernet'});bay(vm,3,'empty')
    assert query(vm)['profiles']['network']==[]
    print('PASS: ENVIRONMENT insertion/eject/pull, PROGRAM explicit unprivileged execution, managed removal, USB Ethernet and on-demand DHCP',flush=True)
with VM(work,'no-dhcp-server',fixture,extra=['-device',controller,'-netdev','hubport,id=isolated,hubid=0','-device','usb-net,id=ethernet,bus=xhci.0,port=3,netdev=isolated']) as vm:
    vm.expect(rb'FDS# ')
    shell(vm,'printf prompt-usable; fds-boottrace mark console-ready','M8_PROMPT_WITHOUT_DHCP')
    wait(vm,'fds --json profiles',lambda s:len(json.loads(s)['profiles']['network'])==1)
    shell(vm,'s6-svwait -u -t 5000 /run/service/dhcp && test -z "$(ip -o -4 address show scope global)" && printf still-usable','M8_DHCP_NOT_A_BARRIER')
    shell(vm,'fds network off','M8_NO_DHCP_STOPPED')
    print('PASS: console and service controls work with Ethernet present and no DHCP server',flush=True)
# Convert this test's raw X11 screenshot to a portable PNG without external tools.
data=(work/'control.xwd').read_bytes();header=struct.unpack('>25I',data[:100]);size,version,fmt,depth,width,height,xoff,order,unit,bitorder,pad,bpp,stride,visual,rm,gm,bm,*_=header
assert version==7 and fmt==2 and bpp==32 and (rm,gm,bm)==(0xff0000,0xff00,0xff)
offset=size+header[19]*12
raw=bytearray()
for y in range(height):
    raw.append(0)
    for x in range(width):
        pixel=int.from_bytes(data[offset+y*stride+x*4:offset+y*stride+x*4+4], 'little' if order==0 else 'big')
        raw.extend(((pixel>>16)&255,(pixel>>8)&255,pixel&255))
def chunk(kind,body):return struct.pack('>I',len(body))+kind+body+struct.pack('>I',zlib.crc32(kind+body)&0xffffffff)
(work/'control.png').write_bytes(b'\x89PNG\r\n\x1a\n'+chunk(b'IHDR',struct.pack('>IIBBBBB',width,height,8,2,0,0,0))+chunk(b'IDAT',zlib.compress(raw))+chunk(b'IEND',b''))
(work/'acceptance.json').write_text(json.dumps(dict(status='passed', native_arm_x11=True, panel_uid=1000, global_grayscale_defaults=True, legacy_and_void_programs_launched_by_mouse=True, keyboard_rescan=True, panel_safe_eject=True, desktop_stop_removes_panel=True, physical_display='not tested'),indent=2)+'\n')
link=project/'out/m8-vm-latest';temporary=link.with_suffix('.next');temporary.symlink_to(work.name);temporary.replace(link)
print(f'PASS: M8 ARM desktop and network verification: {work}')
print('SKIP: Pi DRM/VC4 output, physical E-Ink quality and monitor input-to-refresh latency')