#!/usr/bin/env python3
"""Exercise the packaged machine settings service against disposable virtual NVMe."""
import json
from pathlib import Path
import shlex
import shutil
import subprocess
import sys
import tempfile
import time

project = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(project / 'tools'))
from image_formats import digest, gpt, LINUX_FILESYSTEM
from vm_test import VM

work = Path(tempfile.mkdtemp(prefix='m12-internal.', dir=project / 'out'))
runner = str(project / 'tools/in-image-tools')
controller = 'qemu-xhci,id=xhci,addr=05.0'
placeholder = work / 'placeholder.bin'
placeholder.write_bytes(bytes(4096))
empty = work / 'empty.img'
gpt(empty, [('FDS_TEST', LINUX_FILESYSTEM, placeholder)])

for name in ['build-a', 'build-b']:
    output = work / name
    output.mkdir()
    with (work / (name + '.log')).open('wb') as log:
        subprocess.run([str(project / 'image/build-internal'), '--output-directory', str(output)], check=True, stdout=log, stderr=subprocess.STDOUT)
image = work / 'build-a/internal.img'
layout = json.loads((work / 'build-a/layout.json').read_text())
assert digest(image) == digest(work / 'build-b/internal.img'), 'Internal image construction is not reproducible'
assert [p['name'] for p in layout['partitions']] == ['FDS_BOOT', 'FDS_RECOVERY', 'FDS_INTERNAL']
assert [p['size'] * 512 for p in layout['partitions']] == [512 * 1024**2, 1024**3, 256 * 1024**2]
assert not ((project / 'out/rootfs-cli.tar').resolve().parent / 'root/var/cache/ldconfig/aux-cache').exists()
print('PASS: two complete internal GPT builds are byte-identical, with independent FAT/EROFS/ext4/GPT checks', flush=True)

# Every VM boots the packaged recovery EROFS. No injected runtime or bay-map
# replacement is used; installation exercises the public persistent-settings API.
def extra(path, node='internal', readonly=False):
    return ['-drive', f'file={path},if=none,id={node},format=raw' + (',readonly=on' if readonly else ''),
            '-device', f'nvme,drive={node},serial=FDS-{node.upper()}']

def wait(operation, condition, timeout=60):
    deadline = time.monotonic() + timeout
    while True:
        result = operation()
        if condition(result): return result
        assert time.monotonic() < deadline, result

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

def enter(vm):
    vm.expect(rb'FDS_SYSTEM MEDIA NOT PRESENT')
    vm.send('recovery')
    vm.expect(rb'RECOVERY# ')
    assert vm.capture('id -u') == '0'
    assert vm.capture('stat -c %u:%g:%a /') == '0:0:755'
    wait(lambda: vm.capture('test -S /run/fds/control.sock && echo ready || echo pending'), lambda text: text == 'ready')
    return query(vm, ['machine', 'status'])

def finish(vm):
    vm.send('fds poweroff')
    vm.expect(rb'reboot: Power down')
    assert vm.child.wait(timeout=20) == 0

def unchanged_payloads(path):
    import hashlib
    with path.open('rb') as source:
        for part in layout['partitions'][:2]:
            source.seek(part['start'] * 512)
            checksum = hashlib.sha256()
            remaining = part['payload_bytes']
            while remaining:
                chunk = source.read(min(remaining, 1024 * 1024))
                assert chunk
                remaining -= len(chunk)
                checksum.update(chunk)
            assert checksum.hexdigest() == part['sha256']

def copy_image(name):
    destination = work / (name + '.img')
    subprocess.run(['cp', '--reflink=auto', '--sparse=always', str(image), str(destination)], check=True)
    return destination

machine = copy_image('machine')
original = digest(machine)
with VM(work, 'install-settings', empty, extra=[*extra(machine), '-device', controller, '-device', 'usb-kbd,bus=xhci.0,port=4']) as vm:
    status = enter(vm)
    assert status['source'] == 'internal_nvme' and status['name'] == 'FP-85', status
    assert digest(machine) == original, 'Loading settings changed the disk'
    vm.capture('test -z "$(findmnt -nr -o TARGET | grep /run/fds/machine/internal)"')
    vm.capture('s6-setuidgid fds fds machine store denied /etc/hostname', ok=False)
    vm.capture('s6-setuidgid fds fds machine --load', ok=False)
    devices = wait(lambda: query(vm, ['topology'])['unmapped'], lambda rows: any('03' in row['interfaces'] for row in rows))
    keyboard = next(row for row in devices if '03' in row['interfaces'])
    hub = keyboard['topology'].rsplit('/', 1)[0]
    config = f'[front]\nhub={json.dumps(hub)}\n[front.ports]\n4=4\n'
    catalog = f'[[device]]\nname="LAB KEYBOARD"\nvendor={json.dumps(keyboard["vendor"])}\nproduct={json.dumps(keyboard["product"])}\n'
    vm.capture('fds machine export /tmp/new-machine')
    vm.capture('fds machine export /tmp/new-machine', ok=False)
    for file, data in [('bays.toml', config), ('hardware-catalog.toml', catalog), ('machine.toml', 'format=1\nname="FP-85 LAB"\n')]:
        vm.capture('printf %s ' + shlex.quote(data) + ' >/tmp/new-machine/' + file)
    vm.capture('fds machine validate /tmp/new-machine && fds machine install /tmp/new-machine')
    assert query(vm, ['machine', 'status'])['name'] == 'FP-85'
    vm.capture('s6-rc -l /run/s6-rc -d change cartridged && s6-rc -l /run/s6-rc -u change cartridged')
    assert query(vm, ['topology'])['unmapped'], 'Machine settings changed underneath the current boot'
    vm.capture('fds --json boot-profile >/tmp/boot.json && fds machine store acceptance-boot.json /tmp/boot.json && fds machine fetch acceptance-boot.json /tmp/retrieved.json && cmp /tmp/boot.json /tmp/retrieved.json')
    vm.capture('fds machine store acceptance-boot.json /tmp/boot.json', ok=False)
    vm.capture('fds machine store ../escape /tmp/boot.json', ok=False)
    vm.capture('fds machine fetch acceptance-boot.json /tmp/retrieved.json', ok=False)
    vm.capture('truncate -s 16777217 /tmp/oversize && fds machine store oversize /tmp/oversize', ok=False)
    vm.capture('test -z "$(findmnt -nr -o TARGET | grep /run/fds/machine/internal)"')
    finish(vm)
unchanged_payloads(machine)
print('PASS: settings load without writes, root-only atomic installation, fixed boot snapshot, bounded diagnostics and exact retrieval; boot/recovery partitions remain unchanged', flush=True)

with VM(work, 'persistent-settings', empty, extra=[*extra(machine), '-device', controller, '-device', 'usb-kbd,bus=xhci.0,port=4']) as vm:
    status = enter(vm)
    assert status['source'] == 'internal_nvme' and status['name'] == 'FP-85 LAB', status
    bay = wait(lambda: query(vm, ['bay', '4'])['bays'][0], lambda row: row['state'] == 'hardware')
    assert bay['name'] == 'LAB KEYBOARD', bay
    vm.capture('fds machine fetch acceptance-boot.json /tmp/saved-boot.json')
    saved = json.loads(vm.capture('cat /tmp/saved-boot.json'))
    assert saved['boot_id'] != query(vm, ['boot-profile'])['boot_id']
    (work / 'persistent-status.json').write_text(json.dumps(status, indent=2) + '\n')
    finish(vm)
print('PASS: independent reboot loads saved bay/catalog/name settings and retrieves the previous boot record', flush=True)

# Independent filesystem inspection after real guest writes.
part = layout['partitions'][2]
settings = work / 'after-guest.ext4'
with machine.open('rb') as source, settings.open('wb') as output:
    source.seek(part['start'] * 512)
    remaining = part['payload_bytes']
    while remaining:
        chunk = source.read(min(remaining, 1024 * 1024))
        assert chunk
        output.write(chunk)
        remaining -= len(chunk)
with (work / 'after-guest-fsck.log').open('wb') as log:
    subprocess.run([runner, 'e2fsck', '-f', '-n', str(settings)], check=True, stdout=log, stderr=subprocess.STDOUT)
previous = subprocess.check_output([runner, 'debugfs', '-R', 'cat /config/previous.json', str(settings)], stderr=subprocess.DEVNULL)
assert json.loads(previous)['name'] == 'FP-85'

# Both recovery and settings are present on USB, but only recovery may be selected
# by its GPT label. USB cannot supply persistent machine settings.
with VM(work, 'usb-lookalike', image, system_usb=True) as vm:
    status = enter(vm)
    assert status['source'] == 'image_defaults' and 'No complete internal NVMe' in status['error'], status
    finish(vm)

# Boot the ordinary CLI SYSTEM with two eligible NVMes already present. Stage0
# has one SYSTEM to choose; the machine settings loader must reject ambiguity.
with VM(work, 'ambiguous-internal', project / 'out/fds-system-cli.img', extra=[*extra(image, 'first', True), *extra(image, 'second', True)]) as vm:
    vm.expect(rb'FDS> ')
    assert vm.capture('id -u') == '1000'
    assert vm.capture('stat -c %u:%g:%a /') == '0:0:755'
    wait(lambda: vm.capture('test -S /run/fds/control.sock && echo ready || echo pending'), lambda text: text == 'ready')
    status = query(vm, ['machine', 'status'])
    assert status['source'] == 'image_defaults' and 'Multiple internal NVMe' in status['error'], status
    finish(vm)
print('PASS: USB settings spoof and two eligible NVMes fall back with explicit diagnostics', flush=True)

# Corrupt configuration and an unclean ext4 are distinct failure cases. Build
# fixtures by changing only the settings partition, retaining the packaged OS.
def altered(name, commands):
    fs = work / (name + '.ext4')
    shutil.copyfile(work / 'build-a/internal.ext4', fs)
    script = work / (name + '.debugfs')
    script.write_text('\n'.join(commands) + '\n')
    subprocess.run([runner, 'debugfs', '-w', '-f', str(script), str(fs)], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    disk = copy_image(name)
    with disk.open('r+b') as output, fs.open('rb') as source:
        output.seek(part['start'] * 512)
        shutil.copyfileobj(source, output)
    return disk
bad_config = work / 'bad.json'
bad_config.write_text('{"format":1,"command":"sh"}\n')
fixtures = [('invalid-settings', altered('invalid-settings', ['rm /config/machine.json', f'write {bad_config} /config/machine.json', 'set_inode_field /config/machine.json uid 0', 'set_inode_field /config/machine.json gid 0']), 'Invalid machine settings'),
            ('writable-settings', altered('writable-settings', ['set_inode_field /config/machine.json mode 0100666']), 'root-owned regular'),
            ('unclean-settings', altered('unclean-settings', ['set_super_value state 0']), 'unclean'),
            ('symlink-settings', altered('symlink-settings', ['rm /config/machine.json', 'symlink /config/machine.json /etc/passwd']), '(os error 40)')]
for name, disk, expected in fixtures:
    before = digest(disk)
    with VM(work, name, empty, extra=extra(disk)) as vm:
        status = enter(vm)
        assert status['source'] == 'image_defaults' and expected in status['error'], status
        vm.capture('test -z "$(findmnt -nr -o TARGET | grep /run/fds/machine/internal)"')
        assert digest(disk) == before, 'Failure path wrote internal storage'
        finish(vm)
print('PASS: invalid JSON, writable/symlink settings and unclean ext4 preserve bytes and leave the recovery console usable', flush=True)
link = project / 'out/m12-internal-latest.next'
link.symlink_to(work.name)
link.replace(project / 'out/m12-internal-latest')
print(f'PASS: internal storage software acceptance: {work}', flush=True)
print('SKIP: Pi EEPROM/NVMe boot, real bay calibration, power-loss and flash durability')