#!/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)
    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)
    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)

# 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(['--image', parts[0][2], '--device', disk, '--file-target', '--dry-run'], 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,
            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: interactive/unattended flash, rejection paths, payload readback and relocated GPT:',work)
print('SKIP: physical devices; fixtures are disposable files with inert filesystem signatures')