#!/usr/bin/env python3
"""Create read-only EROFS and a GPT SYSTEM image from an exported rootfs."""
import argparse
import io
import json
from pathlib import Path
import subprocess
import sys
import tarfile
import tempfile
import uuid

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

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--rootfs', type=Path, default=project/'out/rootfs-aarch64.tar')
parser.add_argument('--profile', choices=['cli', 'development'], default=None)
parser.add_argument('--output-directory', type=Path)
args = parser.parse_args()
rootfs = args.rootfs.resolve(strict=True)
with tarfile.open(rootfs) as archive:
    embedded = archive.extractfile('usr/share/fds/image-profile').read(64).decode().strip()
    try:
        root = archive.getmember('.')
    except KeyError:
        parser.error('rootfs omits root directory metadata; rebuild the rootfs')
    if not root.isdir() or (root.uid, root.gid, root.mode) != (0, 0, 0o755):
        parser.error('rootfs must explicitly record / as root:root mode 0755; rebuild the rootfs')
if embedded not in ('cli', 'development'): parser.error('unsupported rootfs profile')
if args.profile is not None and args.profile != embedded: parser.error('requested profile does not match rootfs contents')
args.profile = embedded
work = args.output_directory or Path(tempfile.mkdtemp(prefix='system-build.', dir=project/'out'))
if not work.is_dir() or any(work.iterdir()): parser.error('output directory must exist and be empty')
epoch = int(subprocess.check_output(['git', '-C', str(project/'vendor/void-packages'), 'show', '-s', '--format=%ct', 'HEAD']))
manifest = f'''format = 1\n[cartridge]\nid = "fds.system.{args.profile}"\nname = "FDS/OS {args.profile.upper()}"\nclass = "system"\nversion = "0.1.0"\n[media]\nwritable = false\n'''.encode()
with tarfile.open(rootfs) as source, tarfile.open(work/'system.tar', 'w', format=tarfile.PAX_FORMAT) as output:
    for member in source:
        if member.name in ('FDS', 'FDS/CARTRIDGE.TOML'): continue
        output.addfile(member, source.extractfile(member) if member.isfile() else None)
    directory = tarfile.TarInfo('FDS')
    directory.type = tarfile.DIRTYPE
    directory.mode = 0o755
    directory.mtime = epoch
    output.addfile(directory)
    info = tarfile.TarInfo('FDS/CARTRIDGE.TOML')
    info.mode = 0o644
    info.mtime = epoch
    info.size = len(manifest)
    output.addfile(info, io.BytesIO(manifest))
image_uuid = uuid.uuid5(NAMESPACE, digest(work/'system.tar'))
subprocess.run([str(project/'tools/in-image-tools'), 'mkfs.erofs', '-T', str(epoch), '-U', str(image_uuid),
                '-b4096', '-L', 'FDS_SYSTEM', '--tar=f', str(work/'system.erofs'), str(work/'system.tar')], check=True)
subprocess.run([str(project/'tools/in-image-tools'), 'fsck.erofs', '--extract', str(work/'system.erofs')], check=True)
layout = gpt(work/'system.img', [('FDS_SYSTEM', LINUX_FILESYSTEM, work/'system.erofs')])
layout['rootfs_sha256'] = digest(rootfs)
layout['image_sha256'] = digest(work/'system.img')
layout['profile'] = args.profile
(work/'layout.json').write_text(json.dumps(layout, indent=2)+'\n')
# An independent implementation checks the actual on-disk partition table.
observed = json.loads(subprocess.check_output(['sfdisk', '--json', str(work/'system.img')]))['partitiontable']
assert observed['label'] == 'gpt' and observed['partitions'][0]['name'] == 'FDS_SYSTEM'
assert observed['partitions'][0]['start'] == layout['partitions'][0]['start']
if args.output_directory is None:
    target = project/'out'/f'fds-system-{args.profile}.img'
    temporary = target.with_suffix('.img.next')
    temporary.symlink_to(work.name+'/system.img')
    temporary.replace(target)
print(f'PASS: GPT FDS_SYSTEM and verified EROFS: {work}')