#!/usr/bin/env python3
"""Assemble and read back a Pi 5 FAT32 boot partition; never access a host disk."""
import argparse
import datetime
import json
import mmap
import os
from pathlib import Path
import shutil
import struct
import subprocess
import sys
import tarfile
import tempfile

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

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--mode', choices=['production', 'development'], default='production')
args = parser.parse_args()
subprocess.run([str(project/'tools/prepare-pi-firmware')], check=True)
subprocess.run([str(project/'tools/prepare-image-tools')], check=True)
epoch = int(subprocess.check_output(['git', '-C', str(project/'vendor/void-packages'), 'show', '-s', '--format=%ct', 'HEAD']))
firmware_config = dict(line.split('=', 1) for line in (project/'config/pi-firmware.conf').read_text().splitlines() if line and not line.startswith('#'))
package = project/'out/cache/boot'/f'rpi-firmware-{firmware_config["PI_FIRMWARE_VERSION"]}.aarch64.xbps'
work = Path(tempfile.mkdtemp(prefix='boot-build.', dir=project/'out'))
files = work/'files'
files.mkdir()
with tarfile.open(package) as archive:
    for member in archive:
        name = member.name.removeprefix('./')
        if not member.isfile(): continue
        relative = Path(name)
        if relative.parent == Path('boot') and relative.suffix in ('.elf', '.dat', '.bin'):
            (files/relative.name).write_bytes(archive.extractfile(member).read())
        elif relative.parent == Path('usr/share/licenses/rpi-firmware'):
            (files/relative.name).write_bytes(archive.extractfile(member).read())
kernel = project/'out/kernel/boot'
shutil.copy2(kernel/'kernel_2712.img', files)
for dtb in sorted(kernel.glob('bcm2712-*.dtb')): shutil.copy2(dtb, files)
shutil.copytree(kernel/'overlays', files/'overlays')
shutil.copy2(project/'out/fds-initramfs.img', files/'fds-initramfs.img')
shutil.copy2(project/'image/pi5/config.txt', files)
shutil.copy2(project/'image/pi5'/f'cmdline-{args.mode}.txt', files/'cmdline.txt')
assert (files/'bcm2712-rpi-5-b.dtb').is_file()
assert len((files/'cmdline.txt').read_text().splitlines()) == 1
for path in sorted(files.rglob('*')): os.utime(path, (epoch, epoch))
image = work/'boot.fat'
runner = str(project/'tools/in-image-tools')
subprocess.run([runner, 'mkfs.fat', '--invariant', '-C', '-F', '32', '-n', 'FDS_BOOT', str(image), '524288'], check=True)
for path in sorted(files.iterdir()):
    subprocess.run([runner, 'mcopy', '-m', '-s', '-i', str(image), str(path), '::/'], check=True)

# mtools preserves file times but creates directory records at wall-clock time.
# Normalize FAT32 directory timestamps without touching file data or long names.
stamp = datetime.datetime.fromtimestamp(epoch, datetime.timezone.utc)
date = ((stamp.year-1980)<<9) | (stamp.month<<5) | stamp.day
clock = (stamp.hour<<11) | (stamp.minute<<5) | (stamp.second//2)
with image.open('r+b') as stream, mmap.mmap(stream.fileno(), 0) as disk:
    sector = struct.unpack_from('<H', disk, 11)[0]
    cluster_bytes = sector*disk[13]
    reserved = struct.unpack_from('<H', disk, 14)[0]
    fat_sectors = struct.unpack_from('<I', disk, 36)[0]
    fat_offset = reserved*sector
    data_offset = (reserved+disk[16]*fat_sectors)*sector
    root = struct.unpack_from('<I', disk, 44)[0]
    visited = set()
    def directory(cluster):
        while cluster < 0x0ffffff8:
            assert cluster >= 2 and cluster not in visited
            visited.add(cluster)
            start = data_offset+(cluster-2)*cluster_bytes
            assert start+cluster_bytes <= len(disk)
            for offset in range(start, start+cluster_bytes, 32):
                if disk[offset] == 0: break
                if disk[offset] == 0xe5 or disk[offset+11] == 0x0f: continue
                disk[offset+13] = (stamp.second%2)*100
                struct.pack_into('<HHH', disk, offset+14, clock, date, date)
                struct.pack_into('<HH', disk, offset+22, clock, date)
                if disk[offset+11] & 0x10 and disk[offset] != ord('.'):
                    child = (struct.unpack_from('<H', disk, offset+20)[0]<<16) | struct.unpack_from('<H', disk, offset+26)[0]
                    directory(child)
            cluster = struct.unpack_from('<I', disk, fat_offset+cluster*4)[0] & 0x0fffffff
    directory(root)
    disk.flush()
subprocess.run([runner, 'fsck.fat', '-n', str(image)], check=True)
readback = work/'readback'
readback.mkdir()
subprocess.run([runner, 'mcopy', '-s', '-i', str(image), '::*', str(readback)], check=True)
expected = {str(path.relative_to(files)): digest(path) for path in sorted(files.rglob('*')) if path.is_file()}
observed = {str(path.relative_to(readback)): digest(path) for path in sorted(readback.rglob('*')) if path.is_file()}
assert observed == expected, 'FAT read-back mismatch'
(work/'manifest.json').write_text(json.dumps({'mode': args.mode, 'bytes': image.stat().st_size,
    'image_sha256': digest(image), 'firmware_sha256': digest(package), 'files': expected}, indent=2)+'\n')
for name, source in [('fds-boot.img', work.name+'/boot.fat'), ('boot-volume', work.name)]:
    target = project/'out'/(name+'.next')
    target.symlink_to(source)
    target.replace(project/'out'/name)
print(f'PASS: 512 MiB FDS_BOOT FAT32 partition image, all files verified: {work}')
print('SKIP: firmware boot, NVMe and display mode validation require a physical Pi')