#!/usr/bin/env python3
"""Exercise packaged machine settings against disposable native SD and 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 sd_extra(path):
    return ['-drive', f'file={path},if=none,id=internal_sd,format=raw',
            '-device', 'sdhci-pci,id=sdhci',
            '-device', 'sd-card,drive=internal_sd']

def sd_capacity(path):
    # QEMU SD cards require a power-of-two capacity. Sparse extension and GPT
    # relocation model flashing the small image onto a larger native card.
    with path.open('r+b') as output:
        output.truncate(2 * 1024**3)
    subprocess.run(['sfdisk', '--relocate', 'gpt-bak-std', str(path)], check=True,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    subprocess.run(['sfdisk', '--verify', str(path)], check=True,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return path

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)

sd = sd_capacity(copy_image('native-sd'))
original_sd = digest(sd)
with VM(work, 'sd-install-settings', empty, extra=sd_extra(sd)) as vm:
    status = enter(vm)
    assert status['source'] == 'internal_sd' and status['name'] == 'FP-85', status
    assert digest(sd) == original_sd, 'Loading SD settings changed the card'
    assert vm.capture('cat /sys/class/block/mmcblk0/device/type') == 'SD'
    vm.capture('fds machine export /tmp/sd-machine')
    vm.capture("printf 'format=1\\nname=\"FP-85 SD\"\\n' >/tmp/sd-machine/machine.toml")
    vm.capture('fds machine validate /tmp/sd-machine && fds machine install /tmp/sd-machine')
    assert query(vm, ['machine', 'status'])['name'] == 'FP-85'
    vm.capture('fds --json boot-profile >/tmp/boot.json && fds machine store sd-boot.json /tmp/boot.json')
    vm.capture('test -z "$(findmnt -nr -o TARGET | grep /run/fds/machine/internal)"')
    finish(vm)
unchanged_payloads(sd)
with VM(work, 'sd-persistent-settings', empty, extra=sd_extra(sd)) as vm:
    status = enter(vm)
    assert status['source'] == 'internal_sd' and status['name'] == 'FP-85 SD', status
    vm.capture('fds machine fetch sd-boot.json /tmp/saved.json')
    assert json.loads(vm.capture('cat /tmp/saved.json'))['boot_id'] != query(vm, ['boot-profile'])['boot_id']
    (work / 'sd-persistent-status.json').write_text(json.dumps(status, indent=2) + '\n')
    finish(vm)

# SD supplies machine storage while the actual normal SYSTEM remains on USB.
with VM(work, 'sd-usb-system', project / 'out/fds-system-cli.img', system_usb=True, extra=sd_extra(sd)) as vm:
    vm.expect(rb'FDS> ')
    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'] == 'internal_sd' and status['name'] == 'FP-85 SD', status
    assert vm.capture('findmnt -no FSTYPE /') == 'erofs'
    finish(vm)
print('PASS: native MMC SD recovery, read-only settings load, explicit writes, reboot persistence and normal USB SYSTEM handoff', flush=True)

# Independent filesystem inspection after real guest writes.
part = layout['partitions'][2]
for source_disk, name in [(machine, 'nvme'), (sd, 'sd')]:
    settings = work / (name + '-after-guest.ext4')
    with source_disk.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 / (name + '-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'
print('PASS: independent ext4 checks after guest SD/NVMe writes and previous-settings retention', flush=True)

# 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 SD/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 SD/NVMe' in status['error'], status
    finish(vm)
print('PASS: USB settings spoof and two eligible NVMes fall back with explicit diagnostics', flush=True)

with VM(work, 'ambiguous-sd-nvme', project / 'out/fds-system-cli.img',
        extra=[*sd_extra(sd), *extra(image, 'old_nvme', True)]) as vm:
    vm.expect(rb'FDS> ')
    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 SD/NVMe' in status['error'], status
    finish(vm)
print('PASS: concurrent eligible SD and NVMe are rejected instead of silently selecting settings', 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)
dirty_sd = sd_capacity(altered('unclean-sd', ['set_super_value state 0']))
before = digest(dirty_sd)
with VM(work, 'unclean-sd', empty, extra=sd_extra(dirty_sd)) as vm:
    status = enter(vm)
    assert status['source'] == 'image_defaults' and 'unclean' in status['error'], status
    assert digest(dirty_sd) == before, 'Unclean SD was modified during fallback'
    finish(vm)
print('PASS: unclean SD falls back without journal replay or persistent writes', 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: physical Pi EEPROM/SD/NVMe boot, maintenance USB boot, real bay calibration, power-loss and flash durability')