#!/usr/bin/env python3
"""Create a complete internal SD/NVMe GPT image as a new ordinary file."""
import argparse
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
import uuid

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

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--boot', type=Path, default=project / 'out/fds-boot.img')
parser.add_argument('--recovery', type=Path, default=project / 'out/fds-recovery.img')
parser.add_argument('--machine-config', type=Path, default=project / 'config/machine')
parser.add_argument('--internal-mib', type=int, default=256)
parser.add_argument('--recovery-mib', type=int, choices=[1024, 2048], default=1024)
parser.add_argument('--output-directory', type=Path)
args = parser.parse_args()
if not 32 <= args.internal_mib <= 65536:
    parser.error('internal settings partition must be 32..65536 MiB')
boot, recovery = args.boot.resolve(strict=True), args.recovery.resolve(strict=True)
if not boot.is_file() or boot.stat().st_size != 512 * 1024 * 1024:
    parser.error('boot input must be a 512 MiB regular FAT32 partition image')
if not recovery.is_file() or not 4096 <= recovery.stat().st_size <= args.recovery_mib * 1024 * 1024:
    parser.error('recovery input must be a regular EROFS image fitting its partition')
with boot.open('rb') as stream:
    header = stream.read(512)
    if header[82:90] != b'FAT32   ' or header[71:82] != b'FDS_BOOT   ' or header[510:] != b'\x55\xaa':
        parser.error('boot input is not FDS_BOOT FAT32')
with recovery.open('rb') as stream:
    stream.seek(1024)
    header = stream.read(128)
    if header[:4] != bytes.fromhex('e2e1f5e0') or header[64:80].rstrip(b'\0') != b'FDS_RECOVERY':
        parser.error('recovery input is not FDS_RECOVERY EROFS')
work = args.output_directory or Path(tempfile.mkdtemp(prefix='internal-build.', dir=project / 'out'))
work = work.resolve(strict=True)
if not work.is_dir() or any(work.iterdir()):
    parser.error('output directory must exist and be empty')
subprocess.run([str(project / 'tools/cargo-build'), '--locked', '--offline', '--release', '--target', 'x86_64-unknown-linux-gnu', '-p', 'fds-cli'], cwd=project, check=True)
host_cli = project / 'target/x86_64-unknown-linux-gnu/release/fds'
tree = work / 'settings'
(tree / 'config').mkdir(parents=True)
(tree / 'diagnostics').mkdir(mode=0o700)
subprocess.run([str(host_cli), 'machine', 'pack', str(args.machine_config.resolve(strict=True)), str(tree / 'config/machine.json')], check=True)
epoch = int(subprocess.check_output(['git', '-C', str(project / 'vendor/void-packages'), 'show', '-s', '--format=%ct', 'HEAD']))
for path in [tree, *sorted(tree.rglob('*'))]:
    os.utime(path, (epoch, epoch))
runner = str(project / 'tools/in-image-tools')
subprocess.run([str(project / 'tools/prepare-image-tools')], check=True)
subprocess.run([runner, 'fsck.fat', '-n', str(boot)], check=True)
subprocess.run([runner, 'fsck.erofs', '--extract', str(recovery)], check=True)
settings = work / 'internal.ext4'
with settings.open('xb') as stream:
    stream.truncate(args.internal_mib * 1024 * 1024)
identity = uuid.uuid5(NAMESPACE, 'internal:' + digest(tree / 'config/machine.json'))
# Source ownership is root inside this unprivileged user namespace, independent
# of the workstation UID. No workstation block devices are exposed.
environment = {**os.environ, 'E2FSPROGS_FAKE_TIME': str(epoch)}
subprocess.run(['bwrap', '--unshare-user', '--uid', '0', '--gid', '0', '--ro-bind', '/', '/',
                '--bind', str(work), str(work), '--dev', '/dev', '--proc', '/proc',
                runner, 'mke2fs', '-q', '-t', 'ext4', '-F', '-b', '4096', '-m', '0',
                '-L', 'FDS_INTERNAL', '-U', str(identity),
                '-E', f'root_owner=0:0,lazy_itable_init=0,lazy_journal_init=0,hash_seed={identity}',
                '-d', str(tree), str(settings)], env=environment, check=True)
# mke2fs copies source ctimes; normalize every populated inode after construction.
commands = work / 'normalize.debugfs'
commands.write_text(''.join(f'set_inode_field {name} {field} @{epoch}\n'
    for name in ['/', '/lost+found', '/config', '/config/machine.json', '/diagnostics']
    for field in ['atime', 'ctime', 'mtime', 'crtime']))
subprocess.run([runner, 'debugfs', '-w', '-f', str(commands), str(settings)], env=environment, check=True)
subprocess.run([runner, 'e2fsck', '-f', '-n', str(settings)], env=environment, check=True)
padded = work / 'recovery-partition.erofs'
shutil.copyfile(recovery, padded)
with padded.open('r+b') as stream:
    stream.truncate(args.recovery_mib * 1024 * 1024)
image = work / 'internal.img'
layout = gpt(image, [('FDS_BOOT', EFI_SYSTEM, boot), ('FDS_RECOVERY', LINUX_FILESYSTEM, padded), ('FDS_INTERNAL', LINUX_FILESYSTEM, settings)])
layout.update(format=1, image_sha256=digest(image), machine_settings_sha256=digest(tree / 'config/machine.json'),
              recovery_payload_sha256=digest(recovery), boot_payload_sha256=digest(boot))
observed = json.loads(subprocess.check_output(['sfdisk', '--json', str(image)]))['partitiontable']
assert observed['label'] == 'gpt' and len(observed['partitions']) == 3
for expected, actual in zip(layout['partitions'], observed['partitions']):
    assert (actual['name'], actual['start'], actual['size']) == (expected['name'], expected['start'], expected['size'])
subprocess.run(['sfdisk', '--verify', str(image)], check=True)
(work / 'layout.json').write_text(json.dumps(layout, indent=2) + '\n')
if args.output_directory is None:
    link = project / 'out/fds-internal.img.next'
    link.symlink_to(work.name + '/internal.img')
    link.replace(project / 'out/fds-internal.img')
print(f'PASS: complete internal SD/NVMe image, three verified GPT payloads: {work}')
print('SKIP: no physical disk was written; Pi firmware/SD boot requires hardware')