#!/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-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')
    (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')
    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')
    started=query(vm,'s6-setuidgid fds fds --json run 4 -- check');assert started['started_pid']>1
    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; fds eject 4','M8_PROGRAM_SAFE')
    assert bay(vm,4,'safe')['consumers']==0;remove(4)
    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/'desktop.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/'desktop.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''))
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')