#!/usr/bin/env python3
"""Twelve real virtual USB disks: boot, concurrent I/O, hotplug and native power."""
import hashlib
import io
import itertools
import json
from pathlib import Path
import re
import shlex
import statistics
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, digest, LINUX_FILESYSTEM
work = Path(tempfile.mkdtemp(prefix='m11-vm.', dir=project/'out'))
inputs=[Path(__file__),project/'tools/vm_test.py',project/'out/rootfs-cli.tar',
        project/'out/kernel/boot/kernel_2712.img',project/'out/fds-initramfs.img']
(work/'inputs.sha256').write_text(''.join(f'{digest(path)}  {path.relative_to(project)}\n' for path in inputs))
(work/'qemu-version.txt').write_bytes(subprocess.check_output([str(project/'tools/in-void'),'qemu-system-aarch64','--version']))
controller = 'qemu-xhci,id=xhci,addr=05.0,p2=12,p3=12'
sequence = itertools.count()

def query(vm, args):
    return json.loads(vm.capture('fds --json '+shlex.join(args)))

def wait(vm, read, predicate, timeout=60):
    deadline = time.monotonic()+timeout
    while True:
        value = read()
        if predicate(value): return value
        assert time.monotonic()<deadline, value

def ready(vm, placement):
    def inventory():
        text=vm.capture('fds --json bays 2>/tmp/m11-bays.err; true')
        if text:return json.loads(text)
        error=vm.capture('cat /tmp/m11-bays.err')
        assert 'Cartridge service unavailable' in error,error
        return None
    report = wait(vm, inventory, lambda r:bool(r) and all(
        b['state']==('mounted_read_write' if placement.get(b['bay'])==1 else 'mounted_read_only' if b['bay'] in placement else 'empty')
        for b in r['bays']))
    assert len(report['bays'])==12
    for b in report['bays']:
        if b['bay'] not in placement: continue
        number=placement[b['bay']]
        assert b['manifest']['cartridge']['id']==f'fds.m11.{number:02}',b
        assert len(b['devices'])==1 and b['devices'][0]['serial']==f'M11-{number:02}',b
        assert b['devices'][0]['topology'].endswith('/'+str(b['bay'])),b
    return report

def finish(vm, reboot=False):
    vm.send('fds '+('reboot' if reboot else 'poweroff'))
    report=json.loads(vm.expect(rb'^FDS_SHUTDOWN_FINAL (\{[^\r\n]*\})\r?\n',timeout=90).group(1))
    assert [e['phase'] for e in report['events']]==[
        'frozen','stopping_desktop','stopping_programs','stopping_network',
        'syncing_and_unmounting_data','unmounting_cartridges','prepared','services_stopped'],report
    ending=vm.expect(rb'\[\s*([0-9]+\.[0-9]+)\] reboot: (Power down|Restarting system)')
    assert (ending.group(2)==b'Restarting system')==reboot
    assert vm.child.wait(timeout=20)==0
    report['virtual_total_ms']=float(ending.group(1))*1000-report['events'][0]['at_ns']/1e6
    (work/f'{vm.name}-shutdown.json').write_text(json.dumps(report,indent=2)+'\n')

# Calibrate from live descriptors; never infer bay identity from /dev/sdX.
with VM(work,'probe',project/'out/fds-system-cli.img',extra=['-device',controller,'-device','usb-kbd,bus=xhci.0,port=12']) as vm:
    vm.expect(rb'FDS> ')
    report=wait(vm,lambda:query(vm,['topology']),lambda r:any('03' in d['interfaces'] for d in r['unmapped']))
    device=next(d for d in report['unmapped'] if '03' in d['interfaces'])
    assert device['topology'].endswith('/12'),device
    hub=device['topology'].rsplit('/',1)[0]
    (work/'probe-topology.json').write_text(json.dumps(report,indent=2)+'\n')
    finish(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))

def fixture(name, root_console=False, gated=False):
    replacements={'etc/fds/bays.toml':config.encode()}
    if root_console:
        replacements['usr/libexec/fds/console-session']=b'#!/bin/bash\n# Isolated test image only.\nexec env HOME=/root bash --login\n'
    additions={'usr/libexec/fds/m11-stress':(project/'out/m11-stress').read_bytes(),
               'usr/libexec/fds/m11-writer':(project/'out/m7-writer').read_bytes(),
               'usr/libexec/fds/m11-capture-hardware':(project/'tools/capture-hardware').read_bytes()}
    if gated:
        with tarfile.open(project/'out/rootfs-cli.tar') as original:
            additions['usr/libexec/fds/m11-daemon-real']=original.extractfile('usr/bin/fds-cartridged').read()
        replacements['usr/bin/fds-cartridged']=b'#!/bin/bash\nset -euo pipefail\nmkfifo -m 0666 /run/m11-gate\nprintf "M11_DAEMON_WAIT\\n" >/dev/console\nIFS= read -r release </run/m11-gate\nexec /usr/libexec/fds/m11-daemon-real "$@"\n'
    with tarfile.open(project/'out/rootfs-cli.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'cli'
        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()
        for path,data in {**replacements,**additions}.items():
            member=tarfile.TarInfo(path);member.size=len(data);member.mode=0o755 if path.startswith('usr/') else 0o644
            dst.addfile(member,io.BytesIO(data))
    destination=work/name;destination.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(destination)],check=True,stdout=log,stderr=subprocess.STDOUT)
    return destination/'system.img'

ordinary=fixture('ordinary')
admin=fixture('admin',True)
media={};layouts={};unchanged={}
for n in range(1,13):
    root=work/f'media{n:02}-root';(root/'FDS').mkdir(parents=True)
    kind='data' if n==1 else 'program'
    (root/'FDS/CARTRIDGE.TOML').write_text(f'format=1\n[cartridge]\nid="fds.m11.{n:02}"\nname="M11 BAY {n:02}"\nclass="{kind}"\nversion="1"\n[media]\nwritable={"true" if n==1 else "false"}\n')
    fs=work/f'media{n:02}.{ "ext4" if n==1 else "erofs"}'
    if n==1:
        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)
    else:
        (root/'payload.bin').write_bytes(bytes([n])*(8*1024*1024))
        subprocess.run([str(project/'tools/in-image-tools'),'mkfs.erofs','--quiet','-b','4096','-T','0',str(fs),str(root)],check=True)
    disk=work/f'media{n:02}.img';layouts[n]=gpt(disk,[(f'FDS_{kind.upper()}',LINUX_FILESYSTEM,fs)])
    media[n]=disk
    if n!=1:unchanged[disk]=digest(disk)

identity={n:n for n in range(1,13)}
def cold_args(placement):
    args=['-device',controller]
    for port,n in placement.items():
        args+=['-drive',f'file={media[n]},if=none,id=media{port},format=raw'+(',readonly=on' if n!=1 else ''),
               '-device',f'usb-storage,id=bay{port},drive=media{port},bus=xhci.0,port={port},serial=M11-{n:02}']
    return args

def insert(vm, placement):
    # Pause virtual CPUs only while changing the fixture. All insertions are
    # pending together when execution resumes; no time-based enumeration delay.
    vm.qmp('stop')
    for port,n in placement.items():
        node=f'hot{next(sequence)}'
        vm.qmp('blockdev-add',{'driver':'raw','node-name':node,'read-only':n!=1,'file':{'driver':'file','filename':str(media[n])}})
        vm.qmp('device_add',{'driver':'usb-storage','id':f'bay{port}','drive':node,'bus':'xhci.0','port':str(port),'serial':f'M11-{n:02}'})
    vm.qmp('cont')

def remove_all(vm):
    # DATA has already been safely ejected; read-only PROGRAM pulls may be hot.
    vm.qmp('stop')
    for port in range(1,13):vm.qmp('device_del',{'id':f'bay{port}'})
    vm.qmp('cont')
    ready(vm,{})

def kernel_check(vm, name):
    kernel=vm.capture('dmesg')
    (work/f'{name}-kernel.log').write_text(kernel+'\n')
    errors=[line for line in kernel.splitlines() if re.search(r'(?i)(I/O error|EXT4-fs error|Buffer I/O|reset .*USB device|device descriptor read.*error|BUG:|Call trace:|hung task|host controller not responding|HC died)',line)]
    assert not errors,errors

# Like-for-like real ordinary-user prompt measurements with the same controller.
reports={'empty':[],'twelve':[]}
for iteration in range(3):
    for name,placement in [('empty',{}),('twelve',identity)]:
        with VM(work,f'cold-{name}-{iteration}',ordinary,extra=cold_args(placement)) as vm:
            vm.expect(rb'FDS> ')
            assert vm.capture('id -u')=='1000'
            trace=query(vm,['boot-profile']);reports[name].append(trace)
            (work/f'{vm.name}-boot.json').write_text(json.dumps(trace,indent=2)+'\n')
            (work/f'{vm.name}-bays.json').write_text(json.dumps(ready(vm,placement),indent=2)+'\n')
            finish(vm)
    print(f'PASS: cold boot pair {iteration+1}: empty and twelve USB disks, stable identities and native shutdown',flush=True)
comparison={name:{'samples_ms':[r['durations_ns']['kernel-to-console']/1e6 for r in rows],
                  'median_ms':statistics.median(r['durations_ns']['kernel-to-console']/1e6 for r in rows)} for name,rows in reports.items()}
comparison['change_ms']=comparison['twelve']['median_ms']-comparison['empty']['median_ms']
(work/'boot-comparison.json').write_text(json.dumps(comparison,indent=2)+'\n')
print(f'OBSERVED: twelve-device VM median console change {comparison["change_ms"]:+.3f} ms; physical timing deferred',flush=True)

# A real readiness gate proves that full cartridge enumeration is not a
# prerequisite of the ordinary user's prompt, even with all bays populated.
gated=fixture('gated',gated=True)
with VM(work,'blocked-cartridged-twelve',gated,extra=cold_args(identity)) as vm:
    vm.expect(rb'FDS> ')
    if b'M11_DAEMON_WAIT' not in vm.data:vm.expect(rb'M11_DAEMON_WAIT')
    vm.capture('test "$(id -u)" = 1000 && test ! -e /run/fds/control.sock')
    (work/'gated-boot.json').write_text(json.dumps(query(vm,['boot-profile']),indent=2)+'\n')
    vm.capture('printf "release\\n" >/run/m11-gate')
    ready(vm,identity);finish(vm)
print('PASS: ordinary-user prompt is independent of blocked cartridge readiness with twelve devices present',flush=True)

# The same packaged image and controller are compared under different device
# loads. Keep the actual median sample reports and the required explanation.
explanation=('Controlled twelve-device load adds concurrent kernel USB work and cartridge scanning to '
             'the same two-vCPU VM. Stage0 discovery and the service graph are unchanged; the blocked '
             'twelve-device cartridge test proves console independence. This is measured device-load '
             'overhead, not a Pi timing claim. All raw samples, including variability, are retained.')
for name,rows in reports.items():
    median=sorted(rows,key=lambda r:r['durations_ns']['kernel-to-console'])[len(rows)//2]
    (work/f'{name}-median.json').write_text(json.dumps(median,indent=2)+'\n')
with (work/'boot-review.log').open('wb') as log:
    subprocess.run([str(project/'tools/in-void'),'qemu-aarch64',str(project/'out/fds-boottrace'),'compare',
                    str(work/'empty-median.json'),str(work/'twelve-median.json'),
                    *(['--explain',explanation] if comparison['change_ms']>100 else [])],check=True,stdout=log,stderr=subprocess.STDOUT)

with VM(work,'hotplug-io-reboot',admin,extra=['-device',controller]) as vm:
    vm.expect(rb'FDS# ')
    for iteration in range(3):
        placement={port:((port+iteration-1)%12)+1 for port in reversed(range(1,13))}
        insert(vm,placement)
        report=ready(vm,placement)
        (work/f'hot-{iteration}-bays.json').write_text(json.dumps(report,indent=2)+'\n')
        data_bay=next(port for port,n in placement.items() if n==1)
        # SAFE must not follow a port's previous occupant when devices rotate.
        assert all(b['state']!='safe' for b in report['bays'])
        query(vm,['eject',str(data_bay)])
        remove_all(vm)
        print(f'PASS: simultaneous hotplug cycle {iteration+1}; reverse insertion order and rotated cartridge identities',flush=True)
    insert(vm,identity);ready(vm,identity)
    before=vm.qmp('query-blockstats')
    vm.capture('echo 3 >/proc/sys/vm/drop_caches')
    paths=[f'/run/fds/apps/fds.m11.{n:02}/payload.bin' for n in range(2,13)]
    rows=[json.loads(line) for line in vm.capture('s6-setuidgid fds /usr/libexec/fds/m11-stress '+shlex.join(paths)).splitlines()]
    assert len(rows)==12 and {r['bay'] for r in rows}==set(range(1,13)),rows
    assert max(r['start_ns'] for r in rows)<min(r['end_ns'] for r in rows),'Workloads did not overlap'
    after=vm.qmp('query-blockstats')
    (work/'concurrent-io.json').write_text(json.dumps({'intervals':rows,'before':before,'after':after},indent=2)+'\n')
    def stats(rows):
        result={}
        for row in rows:
            match=re.search(r'(?:^|/)bay([0-9]+)(?:/|$)',row.get('qdev',''))
            if match:result[int(match.group(1))]=row['stats']
        assert len(result)==12,rows
        return result
    old,new=stats(before),stats(after)
    for port in range(1,13):
        key=port
        metric='wr_bytes' if port==1 else 'rd_bytes'
        minimum=64*1024*1024 if port==1 else 32*1024*1024
        assert new[key][metric]-old[key][metric]>=minimum,(port,metric,old[key],new[key])
    kernel_check(vm,'concurrent')
    vm.capture('/bin/bash /usr/libexec/fds/m11-capture-hardware /tmp/m11-capture && cd /tmp/m11-capture && sha256sum -c SHA256SUMS && test "$(stat -c %a .)" = 700')
    (work/'hardware-capture-status.tsv').write_text(vm.capture('cat /tmp/m11-capture/status.tsv')+'\n')
    print('PASS: all twelve I/O intervals overlap; eleven direct readers verify payloads while DATA writes and flushes 64 MiB',flush=True)
    query(vm,['run','1','--','/usr/libexec/fds/m11-writer'])
    wait(vm,lambda:vm.capture('cat /data/progress 2>/dev/null || printf 0'),lambda s:s.isdigit() and int(s)>=128)
    finish(vm,reboot=True)

# Verify DATA independently, then boot the exact disks in another device order.
part=layouts[1]['partitions'][0];payload=work/'data-after.ext4'
with media[1].open('rb') as src,payload.open('wb') as dst:
    src.seek(part['start']*512);left=part['payload_bytes']
    while left:
        chunk=src.read(min(left,1024*1024));assert chunk;dst.write(chunk);left-=len(chunk)
with (work/'data-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()
with VM(work,'reboot-followup',admin,extra=cold_args(dict(reversed(list(identity.items()))))) as vm:
    vm.expect(rb'FDS# ');ready(vm,identity)
    assert vm.capture('stat -c %s /data/stress.bin')==str(64*1024*1024)
    kernel_check(vm,'reboot-followup')
    finish(vm)
for path,checksum in unchanged.items():assert digest(path)==checksum,path
link=project/'out/m11-vm-latest';link.unlink(missing_ok=True);link.symlink_to(work.name)
print(f'PASS: M11 twelve-device virtual stress: {work}',flush=True)
print('SKIP: physical hub wiring, USB voltage/current/reset behavior, RP1, battery and Pi timing')