#!/usr/bin/env python3
"""Format a real FAT32 image from a set of host files — the danos boot volume.

Mirrors tools/make-initial-ramdisk.py in spirit: pure Python 3 standard library,
no external tools (no mkfs.fat / mtools). It writes a valid FAT32 filesystem — a
boot sector + BPB, an FSInfo sector, a backup boot sector, two FATs, and a
directory tree of clusters — so UEFI/OVMF boots \\EFI\\BOOT\\BOOTX64.efi off it
and the danos FAT driver mounts the same image.

    make-fat-image.py <out.img> <size-MiB> [<dest-path> <host-file>]...
    make-fat-image.py --verify <out.img>

Each <dest-path> is a forward-slash path inside the image (e.g.
"EFI/BOOT/BOOTX64.efi"); intermediate directories are created. Names that do not
fit 8.3 get a mangled short name plus long-file-name (LFN) entries.
"""

import struct
import sys

SECTOR = 512
SECTORS_PER_CLUSTER = 1  # 512-byte clusters keep the cluster count high for FAT32
RESERVED_SECTORS = 32
NUM_FATS = 2
CLUSTER_BYTES = SECTOR * SECTORS_PER_CLUSTER

END_OF_CHAIN = 0x0FFFFFFF
BAD_CLUSTER = 0x0FFFFFF7

ATTR_ARCHIVE = 0x20
ATTR_DIRECTORY = 0x10
ATTR_LONG_NAME = 0x0F

VALID_83 = set("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789$%'-_@~!(){}^#& ")


def fat32_geometry(total_sectors):
    """Solve for the FAT size (sectors per FAT) and cluster count that fit."""
    fat_size = 1
    while True:
        data_sectors = total_sectors - RESERVED_SECTORS - NUM_FATS * fat_size
        cluster_count = data_sectors // SECTORS_PER_CLUSTER
        needed = ((cluster_count + 2) * 4 + SECTOR - 1) // SECTOR
        if needed <= fat_size:
            return fat_size, cluster_count
        fat_size = needed


class Fat32Image:
    def __init__(self, total_sectors):
        self.total_sectors = total_sectors
        self.fat_size, self.cluster_count = fat32_geometry(total_sectors)
        if self.cluster_count < 65525:
            sys.exit(f"error: image too small for FAT32 ({self.cluster_count} clusters "
                     f"< 65525); use a larger size")
        self.first_data_sector = RESERVED_SECTORS + NUM_FATS * self.fat_size
        # The FAT, in memory: entry 0 media, entry 1 EOC, entry 2 the root dir.
        self.fat = [0] * (self.cluster_count + 2)
        self.fat[0] = 0x0FFFFFF8
        self.fat[1] = END_OF_CHAIN
        self.fat[2] = END_OF_CHAIN
        self.next_free = 3
        self.cluster_data = {}  # cluster number -> bytes (one cluster's worth)

    def alloc(self):
        cluster = self.next_free
        if cluster >= self.cluster_count + 2:
            sys.exit("error: image out of clusters")
        self.next_free += 1
        self.fat[cluster] = END_OF_CHAIN
        return cluster

    def store_chain(self, content):
        """Allocate a cluster chain holding `content` and return its first cluster."""
        length = max(1, (len(content) + CLUSTER_BYTES - 1) // CLUSTER_BYTES)
        clusters = [self.alloc() for _ in range(length)]
        for i in range(length - 1):
            self.fat[clusters[i]] = clusters[i + 1]
        for i, cluster in enumerate(clusters):
            chunk = content[i * CLUSTER_BYTES:(i + 1) * CLUSTER_BYTES]
            self.cluster_data[cluster] = chunk + b"\x00" * (CLUSTER_BYTES - len(chunk))
        return clusters[0]

    def store_directory(self, first_cluster, entries):
        """Write directory `entries` (bytes) into `first_cluster`, extending the chain."""
        length = max(1, (len(entries) + CLUSTER_BYTES - 1) // CLUSTER_BYTES)
        clusters = [first_cluster]
        for _ in range(length - 1):
            clusters.append(self.alloc())
        for i in range(len(clusters) - 1):
            self.fat[clusters[i]] = clusters[i + 1]
        for i, cluster in enumerate(clusters):
            chunk = entries[i * CLUSTER_BYTES:(i + 1) * CLUSTER_BYTES]
            self.cluster_data[cluster] = chunk + b"\x00" * (CLUSTER_BYTES - len(chunk))

    def cluster_sector(self, cluster):
        return self.first_data_sector + (cluster - 2) * SECTORS_PER_CLUSTER

    def serialize(self):
        image = bytearray(self.total_sectors * SECTOR)
        image[0:SECTOR] = self.boot_sector()
        image[SECTOR:2 * SECTOR] = self.fsinfo_sector()
        image[6 * SECTOR:7 * SECTOR] = self.boot_sector()  # backup boot sector
        # Both FATs.
        fat_bytes = b"".join(struct.pack("<I", entry & 0x0FFFFFFF) for entry in self.fat)
        fat_bytes += b"\x00" * (self.fat_size * SECTOR - len(fat_bytes))
        for copy in range(NUM_FATS):
            base = (RESERVED_SECTORS + copy * self.fat_size) * SECTOR
            image[base:base + len(fat_bytes)] = fat_bytes
        # The data region (clusters).
        for cluster, data in self.cluster_data.items():
            base = self.cluster_sector(cluster) * SECTOR
            image[base:base + len(data)] = data
        return bytes(image)

    def boot_sector(self):
        sector = bytearray(SECTOR)
        # BPB.
        struct.pack_into(
            "<3s8sHBHBHHBHHHII", sector, 0,
            b"\xEB\x58\x90",         # jump
            b"MSWIN4.1",             # OEM name (widest firmware compatibility)
            SECTOR,                  # bytes per sector
            SECTORS_PER_CLUSTER,     # sectors per cluster
            RESERVED_SECTORS,        # reserved sector count
            NUM_FATS,                # number of FATs
            0,                       # root entry count (0 for FAT32)
            0,                       # total sectors 16 (0 -> use 32)
            0xF8,                    # media descriptor
            0,                       # FAT size 16 (0 for FAT32)
            32,                      # sectors per track
            2,                       # heads
            0,                       # hidden sectors
            self.total_sectors,      # total sectors 32
        )
        # FAT32 extended BPB (offset 36).
        struct.pack_into(
            "<IHHIHH12sBBBI11s8s", sector, 36,
            self.fat_size,           # FAT size 32
            0,                       # extended flags
            0,                       # filesystem version
            2,                       # root cluster
            1,                       # FSInfo sector
            6,                       # backup boot sector
            b"\x00" * 12,            # reserved
            0x80,                    # drive number
            0,                       # reserved
            0x29,                    # extended boot signature
            0x12345678,              # volume id
            b"DANOS      ",          # volume label
            b"FAT32   ",             # filesystem type
        )
        sector[510] = 0x55
        sector[511] = 0xAA
        return bytes(sector)

    def fsinfo_sector(self):
        sector = bytearray(SECTOR)
        struct.pack_into("<I", sector, 0, 0x41615252)      # lead signature
        struct.pack_into("<I", sector, 484, 0x61417272)    # struct signature
        free = self.cluster_count - (self.next_free - 2)
        struct.pack_into("<I", sector, 488, free)          # free count
        struct.pack_into("<I", sector, 492, self.next_free)  # next free hint
        struct.pack_into("<I", sector, 508, 0xAA550000)    # trail signature
        return bytes(sector)


def lfn_checksum(short_name):
    checksum = 0
    for byte in short_name:
        checksum = (((checksum & 1) << 7) + (checksum >> 1) + byte) & 0xFF
    return checksum


def short_name_for(name, used):
    """Return (raw 11-byte 8.3 name, needs_lfn)."""
    if "." in name and not name.startswith("."):
        base, ext = name.rsplit(".", 1)
    else:
        base, ext = name, ""
    upper_base, upper_ext = base.upper(), ext.upper()
    # A name fits 8.3 if it is short enough and uses valid characters. The raw
    # 8.3 entry is always uppercase; if that loses the real name's case (e.g.
    # "init" -> "INIT"), a long-name chain carries the exact name. This matters
    # because the EFI loader *enumerates* /system to build the ramdisk — it gets
    # back whatever the directory stores, so the stored name must be exact, not
    # merely case-insensitively findable.
    fits = (1 <= len(base) <= 8 and len(ext) <= 3
            and all(c in VALID_83 for c in upper_base + upper_ext))
    if fits:
        exact = base == upper_base and ext == upper_ext
        return (upper_base.ljust(8) + upper_ext.ljust(3)).encode("ascii"), not exact
    # Mangle to STEM~N.EXT.
    stem = "".join(c for c in upper_base if c in VALID_83 and c != " ")[:6] or "FILE"
    index = 1
    while True:
        candidate = f"{stem}~{index}".ljust(8)[:8] + upper_ext.ljust(3)[:3]
        raw = candidate.encode("ascii")
        if raw not in used:
            used.add(raw)
            return raw, True
        index += 1


def lfn_entries(name, short_raw):
    checksum = lfn_checksum(short_raw)
    units = list(name.encode("utf-16-le"))
    pairs = [bytes(units[i:i + 2]) for i in range(0, len(units), 2)]
    pairs.append(b"\x00\x00")  # null terminator
    while len(pairs) % 13 != 0:
        pairs.append(b"\xff\xff")
    count = len(pairs) // 13
    out = bytearray()
    for sequence in range(count, 0, -1):  # stored last-logical-first
        piece = pairs[(sequence - 1) * 13:sequence * 13]
        entry = bytearray(32)
        entry[0] = sequence | (0x40 if sequence == count else 0)
        for i in range(5):
            entry[1 + i * 2:1 + i * 2 + 2] = piece[i]
        entry[11] = ATTR_LONG_NAME
        entry[12] = 0
        entry[13] = checksum
        for i in range(6):
            entry[14 + i * 2:14 + i * 2 + 2] = piece[5 + i]
        entry[26:28] = b"\x00\x00"
        for i in range(2):
            entry[28 + i * 2:28 + i * 2 + 2] = piece[11 + i]
        out += entry
    return bytes(out)


def short_entry(raw11, attributes, cluster, size):
    return struct.pack(
        "<11sBBBHHHHHHHI",
        raw11, attributes, 0, 0, 0, 0, 0,
        (cluster >> 16) & 0xFFFF, 0, 0, cluster & 0xFFFF, size,
    )


def write_directory(image, cluster, children, parent_cluster, is_root):
    """Recursively lay out a directory: allocate child clusters, build entries."""
    entries = bytearray()
    if not is_root:
        entries += short_entry(b".          ", ATTR_DIRECTORY, cluster, 0)
        parent = 0 if parent_cluster == 2 else parent_cluster
        entries += short_entry(b"..         ", ATTR_DIRECTORY, parent, 0)
    used_short_names = set()
    for name, child in children.items():
        raw, needs_lfn = short_name_for(name, used_short_names)
        used_short_names.add(raw)
        if child["type"] == "dir":
            child_cluster = image.alloc()
            if needs_lfn:
                entries += lfn_entries(name, raw)
            entries += short_entry(raw, ATTR_DIRECTORY, child_cluster, 0)
            write_directory(image, child_cluster, child["children"], cluster, False)
        else:
            data = child["data"]
            first = image.store_chain(data) if data else 0
            if needs_lfn:
                entries += lfn_entries(name, raw)
            entries += short_entry(raw, ATTR_ARCHIVE, first, len(data))
    image.store_directory(cluster, bytes(entries))


def build_tree(pairs):
    root = {}
    for dest, host in pairs:
        with open(host, "rb") as handle:
            data = handle.read()
        parts = [p for p in dest.replace("\\", "/").split("/") if p]
        node = root
        for part in parts[:-1]:
            node = node.setdefault(part, {"type": "dir", "children": {}})["children"]
        node[parts[-1]] = {"type": "file", "data": data}
    return root


def build(out_path, size_mib, pairs):
    total_sectors = size_mib * 1024 * 1024 // SECTOR
    image = Fat32Image(total_sectors)
    tree = build_tree(pairs)
    write_directory(image, 2, tree, 0, True)
    with open(out_path, "wb") as handle:
        handle.write(image.serialize())
    print(f"make-fat-image: wrote {out_path} "
          f"({size_mib} MiB FAT32, {image.cluster_count} clusters)")


def verify(path):
    with open(path, "rb") as handle:
        data = handle.read()
    if len(data) < SECTOR or data[510] != 0x55 or data[511] != 0xAA:
        sys.exit("verify: missing 0x55AA boot signature")
    bytes_per_sector, sectors_per_cluster = struct.unpack_from("<HB", data, 11)
    reserved, num_fats = struct.unpack_from("<H", data, 14)[0], data[16]
    fat_size_32, root_cluster = struct.unpack_from("<I", data, 36)[0], struct.unpack_from("<I", data, 44)[0]
    total_sectors = struct.unpack_from("<I", data, 32)[0]
    if bytes_per_sector != SECTOR or sectors_per_cluster == 0 or num_fats == 0 or fat_size_32 == 0:
        sys.exit("verify: implausible BPB")
    first_data = reserved + num_fats * fat_size_32
    cluster_count = (total_sectors - first_data) // sectors_per_cluster
    if cluster_count < 65525:
        sys.exit(f"verify: not FAT32 ({cluster_count} clusters)")
    # Resolve EFI/BOOT/BOOTX64.efi through the directory tree to prove it is present.
    if not _resolve(data, ["EFI", "BOOT", "BOOTX64.EFI"], root_cluster,
                    reserved, num_fats, fat_size_32, first_data, sectors_per_cluster):
        sys.exit("verify: EFI/BOOT/BOOTX64.efi not found")
    print(f"verify: {path} is FAT32 ({cluster_count} clusters); EFI/BOOT/BOOTX64.efi present")


def _read_fat(data, cluster, reserved):
    offset = reserved * SECTOR + cluster * 4
    return struct.unpack_from("<I", data, offset)[0] & 0x0FFFFFFF


def _resolve(data, parts, cluster, reserved, num_fats, fat_size, first_data, spc):
    for part in parts:
        cluster = _find(data, cluster, part, reserved, first_data, spc)
        if cluster is None:
            return False
    return True


def _find(data, dir_cluster, name, reserved, first_data, spc):
    target = name.upper()
    cluster = dir_cluster
    guard = 0
    while cluster >= 2 and cluster < BAD_CLUSTER and guard < 100000:
        sector = first_data + (cluster - 2) * spc
        for s in range(spc):
            base = (sector + s) * SECTOR
            for i in range(SECTOR // 32):
                entry = data[base + i * 32:base + i * 32 + 32]
                if entry[0] == 0x00:
                    return None
                if entry[0] == 0xE5 or (entry[11] & ATTR_LONG_NAME) == ATTR_LONG_NAME:
                    continue
                raw = entry[0:11]
                short = (raw[0:8].rstrip().decode("latin1") +
                         ("." + raw[8:11].rstrip().decode("latin1") if raw[8:11].strip() else "")).upper()
                if short == target:
                    return ((entry[20] | (entry[21] << 8)) << 16) | (entry[26] | (entry[27] << 8))
        cluster = _read_fat(data, cluster, reserved)
        guard += 1
    return None


def main(argv):
    if len(argv) == 3 and argv[1] == "--verify":
        verify(argv[2])
        return 0
    if len(argv) < 3 or (len(argv) - 3) % 2 != 0:
        sys.exit("usage: make-fat-image.py <out.img> <size-MiB> [<dest> <host>]...\n"
                 "       make-fat-image.py --verify <out.img>")
    out_path = argv[1]
    size_mib = int(argv[2])
    rest = argv[3:]
    pairs = [(rest[i], rest[i + 1]) for i in range(0, len(rest), 2)]
    build(out_path, size_mib, pairs)
    return 0


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