#!/usr/bin/env python3
"""Build a deterministic Rust-only initramfs with the required early display daemon."""
import argparse
import gzip
import json
from pathlib import Path
import shutil
import stat
import subprocess
import sys
import tempfile

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

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--stage0', type=Path, default=project/'out/fds-stage0')
parser.add_argument('--output-directory', type=Path)
args = parser.parse_args()
payloads = {'fds-stage0': args.stage0.resolve(strict=True), 'dasungd': project/'out/dasungd'}

for tool in ('lz4', 'zstd'):
    if not shutil.which(tool): sys.exit(f'ERROR: missing host {tool}; see fds-os.wiki/Developer-M4-Work.md')
for executable in ('fds-stage0', 'dasungd'):
    subprocess.run([str(project/'tools/verify-elf'), str(payloads[executable]), 'aarch64', 'static'], check=True)
epoch = int(subprocess.check_output(['git', '-C', str(project/'vendor/void-packages'), 'show', '-s', '--format=%ct', 'HEAD']))
work = args.output_directory or Path(tempfile.mkdtemp(prefix='initramfs-build.', dir=project/'out'))
if not work.is_dir() or any(work.iterdir()): parser.error('output directory must exist and be empty')
entries = []
for name in ['dev', 'proc', 'sys', 'newroot', 'sbin', 'etc', 'usr', 'usr/lib', 'usr/lib/firmware', 'usr/lib/firmware/edid']:
    entries.append((name, stat.S_IFDIR | 0o755, b'', 0, 0))
entries.append(('dev/console', stat.S_IFCHR | 0o600, b'', 5, 1))
entries.append(('init', stat.S_IFLNK | 0o777, b'sbin/fds-stage0', 0, 0))
entries.append(('lib', stat.S_IFLNK | 0o777, b'usr/lib', 0, 0))
for executable in ('fds-stage0', 'dasungd'):
    entries.append(('sbin/'+executable, stat.S_IFREG | 0o755, payloads[executable].read_bytes(), 0, 0))
config = (project/'packages/fds-dasungd/files/dasungd.toml').read_text().replace('/run/dasungd/', '/dev/fds-early/')
entries.append(('etc/dasungd-early.toml', stat.S_IFREG | 0o600, config.encode(), 0, 0))
entries.append(('usr/lib/firmware/edid/dasung-paperlike13k-37hz.bin', stat.S_IFREG | 0o644,
                (project/'rust/dasungd/profiles/paperlike13k-37hz.edid').read_bytes(), 0, 0))
archive = newc(sorted(entries), epoch)
(work/'initramfs.cpio').write_bytes(archive)
(work/'initramfs.cpio.gz').write_bytes(gzip.compress(archive, compresslevel=9, mtime=0))
for tool, options, suffix in [('lz4', ['-q', '-l', '-9', '-c'], 'lz4'), ('zstd', ['-q', '-T1', '-19', '-c'], 'zst')]:
    with (work/f'initramfs.cpio.{suffix}').open('wb') as stream:
        subprocess.run([tool, *options], input=archive, stdout=stream, check=True)
# All variants must decode to the exact same cpio stream.
assert gzip.decompress((work/'initramfs.cpio.gz').read_bytes()) == archive
for tool, suffix in [('lz4', 'lz4'), ('zstd', 'zst')]:
    assert subprocess.check_output([tool, '-q', '-d', '-c', str(work/f'initramfs.cpio.{suffix}')]) == archive
manifest = {path.name: {'bytes': path.stat().st_size, 'sha256': digest(path)} for path in sorted(work.glob('initramfs.cpio*'))}
(work/'formats.json').write_text(json.dumps(manifest, indent=2)+'\n')
for target, source in ([] if args.output_directory else [('initramfs', work.name), ('fds-initramfs.img', work.name+'/initramfs.cpio')]):
    temporary = project/'out'/(target+'.next')
    temporary.symlink_to(source)
    temporary.replace(project/'out'/target)
print('PASS: uncompressed, gzip, legacy-LZ4 and Zstandard initramfs variants agree')
print(f'Initramfs artifacts: {work}')
print('SKIP: compression boot-time comparison requires the physical Pi')