#!/usr/bin/env python3
"""Exercise the actual interactive and unattended flasher using disposable files."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import pty
import selectors
import shutil
import struct
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 gpt, LINUX_FILESYSTEM, EFI_SYSTEM, digest

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--cli', type=Path, required=True)
args = parser.parse_args()
cli = args.cli.resolve()
work = Path(tempfile.mkdtemp(prefix='workstation-flash.', dir=project / 'out'))
log = (work / 'commands.log').open('w')

# Deliberately inert filesystem signatures in small independent GPT fixtures.
# These test the writer and table validation, not filesystem usability or boot.
parts = []
for name, kind, signature, offset in [
    ('FDS_BOOT', EFI_SYSTEM, b'FAT32   ', 82),
    ('FDS_RECOVERY', LINUX_FILESYSTEM, bytes.fromhex('e2e1f5e0'), 1024),
    ('FDS_INTERNAL', LINUX_FILESYSTEM, bytes.fromhex('53ef'), 1080),
]:
    path = work / (name + '.bin')
    with path.open('xb') as stream:
        stream.truncate(1024 * 1024)
        stream.seek(offset); stream.write(signature)
        if name == 'FDS_BOOT':
            stream.seek(11); stream.write(b'\0\x02')
            stream.seek(510); stream.write(b'\x55\xaa')
        stream.seek(8192); stream.write(b'known payload, not a bootable filesystem')
    parts.append((name, kind, path))
internal = work / 'internal.img'
gpt(internal, parts)
system = work / 'system.img'
gpt(system, [('FDS_SYSTEM', LINUX_FILESYSTEM, parts[1][2])])


def invoke(arguments, ok=True):
    result = subprocess.run([str(cli), *map(str, arguments)], capture_output=True, text=True, timeout=60)
    log.write(repr(arguments) + '\n' + result.stdout + result.stderr); log.flush()
    assert (result.returncode == 0) == ok, (arguments, result.returncode, result.stdout, result.stderr)
    return result


def target(name, size):
    path = work / name
    with path.open('xb') as stream: stream.truncate(size)
    return path


def preview(image, disk):
    return json.loads(invoke(['--image', image, '--device', disk, '--file-target', '--dry-run', '--json']).stdout)['plan']


def unattended(image, disk, plan, **changes):
    return ['--image', image, '--device', disk, '--file-target', '--unattended',
            '--expect-target', changes.get('target_id', plan['target_id']),
            '--sha256', changes.get('sha256', plan['sha256']), '--json']


def verify(image, disk):
    subprocess.run(['sfdisk', '--verify', str(disk)], check=True, stdout=log, stderr=log)
    original = json.loads(subprocess.check_output(['sfdisk', '--json', str(image)]))['partitiontable']
    written = json.loads(subprocess.check_output(['sfdisk', '--json', str(disk)]))['partitiontable']
    assert original['id'] == written['id']
    assert written['lastlba'] == disk.stat().st_size // 512 - 34
    with image.open('rb') as source, disk.open('rb') as output:
        for a, b in zip(original['partitions'], written['partitions'], strict=True):
            for key in ('start', 'size', 'type', 'uuid', 'name'): assert a[key] == b[key], (a,b)
            source.seek(a['start']*512); output.seek(b['start']*512)
            assert source.read(a['size']*512) == output.read(b['size']*512)
    if disk.stat().st_size == image.stat().st_size: assert digest(image) == digest(disk)


for image in [internal, system]:
    source_hash = digest(image)
    inspection = json.loads(invoke(['inspect', image, '--json']).stdout)
    assert inspection['status'] == 'inspected'
    assert inspection['inspection']['sha256'] == source_hash
    assert inspection['inspection']['image']['kind'] == image.stem
    assert source_hash in invoke(['inspect', image]).stdout
    for extra in [0, 512, 4*1024*1024]:
        disk = target(f'{image.stem}-{extra}.target', image.stat().st_size + extra)
        before = digest(disk)
        plan = preview(image, disk)
        assert digest(disk) == before
        result = invoke(unattended(image, disk, plan, target_id='0'*64), False)
        assert 'Target identity' in result.stderr and digest(disk) == before
        result = invoke(unattended(image, disk, plan, sha256='0'*64), False)
        assert 'SHA-256' in result.stderr and digest(disk) == before
        assert json.loads(invoke(unattended(image, disk, plan)).stdout)['status'] == 'verified'
        verify(image, disk)
        invoke(unattended(image, disk, plan), False)  # A file target's contents changed.
    assert digest(image) == source_hash

for name, partitions in [
    ('data', [('FDS_DATA', LINUX_FILESYSTEM, parts[2][2])]),
    ('environment', [('FDS_ENVIRONMENT', LINUX_FILESYSTEM, parts[1][2])]),
    ('program', [('FDS_PROGRAM', LINUX_FILESYSTEM, parts[1][2])]),
    ('software', [(label, LINUX_FILESYSTEM, parts[1][2]) for label in ('FDS_METADATA','FDS_PAYLOAD02','FDS_PAYLOAD03')]),
]:
    image=work/(name+'.img');gpt(image,partitions)
    inspection=json.loads(invoke(['inspect',image,'--json']).stdout)['inspection']
    assert inspection['sha256']==digest(image)
    assert [p['name'] for p in inspection['image']['partitions']]==[p[0] for p in partitions]
    disk=target(name+'.target',image.stat().st_size+1024*1024)
    plan=preview(image,disk)
    assert json.loads(invoke(unattended(image,disk,plan)).stdout)['status']=='verified'
    verify(image,disk)

# The consolidated command creates a real DATA filesystem, inspects it, and
# flashes it through the same unattended writer used for existing images.
# e2fsprogs is required for this acceptance check, as for DATA creation itself.
tree=work/'data-tree';(tree/'FDS').mkdir(parents=True)
(tree/'FDS/CARTRIDGE.TOML').write_text('format=1\n[cartridge]\nid="flash.data"\nname="Flash test"\nclass="data"\nversion="1"\n[media]\nwritable=true\n')
(tree/'sample.txt').write_text('Created and flashed with fds-flash\n')
created=work/'created-data.img'
result=json.loads(invoke(['create','data',tree,created,'--size-mib','32','--json']).stdout)
assert result['status']=='created'
inspection=json.loads(invoke(['inspect',created,'--json']).stdout)['inspection']
assert result['inspection']==inspection and inspection['sha256']==digest(created)
disk=target('created-data.target',created.stat().st_size+1024*1024)
assert json.loads(invoke(unattended(created,disk,preview(created,disk))).stdout)['status']=='verified'
verify(created,disk)
partition=inspection['image']['partitions'][0]
payload=work/'created-data.ext4'
with disk.open('rb') as stream:
    stream.seek(partition['start']);payload.write_bytes(stream.read(partition['bytes']))
readback=subprocess.run(['debugfs','-R','cat /sample.txt',str(payload)],capture_output=True,text=True,check=True)
assert readback.stdout==(tree/'sample.txt').read_text(),readback
before=digest(created)
invoke(['create','data',tree,created,'--size-mib','32'],False)
assert digest(created)==before
for arguments in [
    ['create','data',tree,tree/'inside.img','--size-mib','32'],
    ['create','environment',tree,work/'wrong-class.img'],
    ['create','system',tree,work/'wrong-system.img'],
    ['create','program',tree,work/'legacy-program.img'],
    ['create','data',tree,work/'bad-size.img','--size-mib','31'],
]:
    invoke(arguments,False)
assert not (tree/'inside.img').exists()
assert not any((work/name).exists() for name in ['wrong-class.img','wrong-system.img','legacy-program.img','bad-size.img'])

# Invalid source geometry and destinations must fail without writing.
disk = target('protected.target', internal.stat().st_size)
before = digest(disk)
invoke(['--image', internal, '--device', disk, '--unattended'], False)
invoke(['--image', internal, '--device', disk, '--file-target'], False)  # No terminal.
invoke(['--image', internal, '--device', disk, '--dry-run'], False)  # Not a block device.
small = target('small.target', 1024)
invoke(['--image', internal, '--device', small, '--file-target', '--dry-run'], False)
invoke(['--image', internal, '--device', internal, '--file-target', '--dry-run'], False)
alias = work/'hardlink.img'; os.link(internal, alias)
invoke(['--image', internal, '--device', alias, '--file-target', '--dry-run'], False)
for name, change in [('bad-primary', 512+16), ('bad-backup', internal.stat().st_size-512+16),
                     ('bad-table', 1024+56), ('bad-filesystem', 2048*512+82)]:
    image = work/(name+'.img'); shutil.copyfile(internal,image)
    with image.open('r+b') as stream:
        stream.seek(change); byte=stream.read(1); stream.seek(change); stream.write(bytes([byte[0]^1]))
    invoke(['--image', image, '--device', disk, '--file-target', '--dry-run'], False)
    invoke(['inspect', image, '--json'], False)
invoke(['--image', parts[0][2], '--device', disk, '--file-target', '--dry-run'], False)
invoke(['inspect', parts[0][2]], False)
assert digest(disk) == before

# Drive the real terminal workflow, including cancellation and changed inputs
# after the review is displayed. No shell evaluates the user's paths or input.
def interactive(name, reply, mutate=None, ok=True, include_image=True):
    disk = target(name+'.target', system.stat().st_size + 1024*1024)
    image = work/(name+'.img'); shutil.copyfile(system,image)
    before = digest(disk)
    master, slave = pty.openpty()
    command = [str(cli), '--device', str(disk), '--file-target', '--json']
    if include_image: command += ['--image', str(image)]
    child = subprocess.Popen(command, stdin=slave, stderr=slave, stdout=subprocess.PIPE, text=True)
    os.close(slave)
    selector=selectors.DefaultSelector();selector.register(master,selectors.EVENT_READ)
    output=bytearray()
    def until(marker):
        deadline=time.monotonic()+20
        while marker not in output:
            assert time.monotonic()<deadline, output
            for _,_ in selector.select(max(0,deadline-time.monotonic())):
                chunk=os.read(master,65536); assert chunk; output.extend(chunk)
    try:
        if not include_image:
            until(b'Image path: ');os.write(master,(str(image)+'\n').encode())
        until(b' to proceed: ')
        if mutate: mutate(image,disk)
        expected = 'ERASE '+str(disk)
        os.write(master,((expected if reply else 'cancel')+'\n').encode())
        deadline=time.monotonic()+30
        while child.poll() is None:
            assert time.monotonic()<deadline, output
            for _,_ in selector.select(1):
                try: output.extend(os.read(master,65536))
                except OSError: pass
        stdout=child.stdout.read()
        assert (child.returncode==0)==ok, (stdout,output)
        if ok:
            assert json.loads(stdout)['status']=='verified'
            verify(image,disk)
        elif mutate is None or name=='changed-source': assert digest(disk)==before
        (work/(name+'.terminal.log')).write_bytes(output+stdout.encode())
    finally:
        if child.poll() is None: child.kill();child.wait()
        child.stdout.close();selector.close();os.close(master)

interactive('interactive-cancel',False,ok=False)
interactive('interactive-write',True,include_image=False)
def change_source(image,disk):
    with image.open('r+b') as stream:stream.seek(1024*1024+8192);stream.write(b'CHANGED')
interactive('changed-source',True,change_source,ok=False)
def change_target(image,disk):
    disk.unlink()
    with disk.open('xb') as stream:stream.truncate(image.stat().st_size+1024*1024)
interactive('replaced-target',True,change_target,ok=False)

actual_images = []
for name in ('fds-internal.img', 'fds-system-cli.img'):
    image=project/'out'/name
    if image.is_file():
        disk=target('actual-'+name+'.target',image.stat().st_size)
        plan=preview(image,disk)
        assert plan['image']['bytes']==image.stat().st_size
        actual_images.append(dict(path=str(image.resolve()),sha256=plan['sha256']))
    else:
        print('SKIP: build image not available for read-only inspection:',image)

record=dict(status='passed',cli=str(cli),cli_sha256=digest(cli),internal_and_all_cartridge_classes=True,interactive_terminal=True,
            image_inspection_without_target=True,data_create_inspect_flash_and_filesystem_readback=True,
            unattended_bindings=True,exact_larger_and_overlapping_gpt=True,independent_sfdisk=True,
            partition_payloads_unchanged=True,invalid_sources_and_targets_rejected=True,
            changed_source_and_target_rejected=True,actual_build_images_inspected=actual_images,physical_disks_written=False)
(work/'acceptance.json').write_text(json.dumps(record,indent=2)+'\n')
(project/'out/workstation-flash-current.txt').write_text(str(work)+'\n')
print('PASS: create/inspect, interactive/unattended flash, rejection paths, payload readback and relocated GPT:',work)
print('SKIP: physical devices; disposable-file fixtures include real DATA and inert signatures for other layouts')