#!/usr/bin/env python3
"""Assemble an MBR-partitioned disk image from N FAT32 partitions — the danos
multi-volume test disk.

Pure Python 3 stdlib (no mtools / parted). Each partition is a real FAT32
filesystem produced by make-fat-image.py, laid out behind a classic MBR so the
danos partition prober (partition.allVolumes) walks the table and the volume
manager spawns one confined filesystem per partition — several volumes sharing
ONE block channel, each clamped to its own LBA range. That shared-channel,
per-partition path is what a single stick with two partitions exercises and a
pair of single-volume sticks does not.

    make-partitioned-image.py <out.img> [<serial-hex> <size-MiB>]...

Each partition is an empty FAT32 with the given volume serial (its /volumes/
fat-<serial> content id). Partitions are 1-MiB aligned; the MBR marks each
type 0x0C (FAT32 LBA). At most four (an MBR holds four primaries).
"""

import os
import struct
import subprocess
import sys
import tempfile

SECTOR = 512
ALIGN = 2048  # sectors (1 MiB) — standard partition alignment, and the gap the MBR sits in
MBR_TYPE_FAT32_LBA = 0x0C
MAX_PRIMARY_PARTITIONS = 4
HERE = os.path.dirname(os.path.abspath(__file__))


def align_up(sectors, to=ALIGN):
    return (sectors + to - 1) // to * to


def main(argv):
    if len(argv) < 4 or (len(argv) - 2) % 2 != 0:
        sys.exit("usage: make-partitioned-image.py <out.img> [<serial-hex> <size-MiB>]...")
    out_path = argv[1]
    specs = [(argv[i], int(argv[i + 1])) for i in range(2, len(argv), 2)]
    if len(specs) > MAX_PRIMARY_PARTITIONS:
        sys.exit(f"error: an MBR holds at most {MAX_PRIMARY_PARTITIONS} primary partitions")

    # Generate each partition's FAT32 image, then place it at its aligned start.
    partitions = []  # (start_sector, sector_count, bytes)
    cursor = ALIGN   # leave the first 1 MiB for the MBR + alignment gap
    with tempfile.TemporaryDirectory() as tmp:
        for idx, (serial, size_mib) in enumerate(specs):
            part_path = os.path.join(tmp, f"p{idx}.img")
            subprocess.run(
                [sys.executable, os.path.join(HERE, "make-fat-image.py"),
                 "--serial", serial, "--label", f"DATA{idx}",
                 part_path, str(size_mib)],
                check=True, stdout=subprocess.DEVNULL)
            with open(part_path, "rb") as handle:
                data = handle.read()
            count = len(data) // SECTOR
            partitions.append((cursor, count, data))
            cursor = align_up(cursor + count)

    total_sectors = cursor
    disk = bytearray(total_sectors * SECTOR)
    # The MBR: a disk signature, one partition entry per FAT partition, 0x55AA.
    # No boot code (this disk is data, never booted); danos's mount() sees the
    # signature but no BPB at LBA 0 and takes the MBR-walk path.
    struct.pack_into("<I", disk, 440, 0x0D05DA05)  # arbitrary but fixed disk signature
    for idx, (start, count, data) in enumerate(partitions):
        entry = 446 + idx * 16
        disk[entry + 0] = 0x00                          # not bootable
        disk[entry + 1:entry + 4] = b"\xFE\xFF\xFF"     # CHS start (LBA-aware tools ignore)
        disk[entry + 4] = MBR_TYPE_FAT32_LBA
        disk[entry + 5:entry + 8] = b"\xFE\xFF\xFF"     # CHS end
        struct.pack_into("<I", disk, entry + 8, start)  # start LBA
        struct.pack_into("<I", disk, entry + 12, count)  # sector count
        disk[start * SECTOR:start * SECTOR + len(data)] = data
    disk[510] = 0x55
    disk[511] = 0xAA

    with open(out_path, "wb") as handle:
        handle.write(disk)
    print(f"make-partitioned-image: wrote {out_path} "
          f"({total_sectors * SECTOR // (1024 * 1024)} MiB, {len(partitions)} partitions)")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))