#!/usr/bin/env python3
"""Create a GPT disk image with FAT32 ESP partition for UEFI boot."""

import struct, sys, os, zlib, uuid

MBR_OFF  = 0
GPT_OFF  = 512
GPT_PARTS_OFF = 1024  # LBA 2
PART_LBA = 34         # partition data starts at LBA 34
SECTOR   = 512

def u16(x): return struct.pack('<H', x)
def u32(x): return struct.pack('<I', x)
def u64(x): return struct.pack('<Q', x)

def crc32(data):
    return struct.pack('<I', zlib.crc32(data) & 0xFFFFFFFF)

def build(efi_path, out_path, mb=64):
    total = mb * 1024 * 1024
    img   = bytearray(total)

    # LBA 0: Protective MBR
    img[0x1FE:0x200] = b'\x55\xAA'          # boot signature
    img[0x1BE:0x1BE+16] = struct.pack('<B3sB3sII',
        0x00,                               # status
        b'\x00\x02\x00',                    # CHS start (fake)
        0xEE,                               # partition type = GPT protective
        b'\xFF\xFF\xFF',                    # CHS end (fake)
        1,                                  # LBA start
        0xFFFFFFFF                          # size (max)
    )

    disk_guid = uuid.uuid4().bytes
    part_guid = uuid.uuid4().bytes
    efi_type  = bytes.fromhex('C12A7328F81F11D2BA4B00A0C93EC93D')
    last_lba  = total // SECTOR - 1

    # Leave last 33 LBAs for backup GPT (32 for partition entries, 1 for header)
    last_part_lba = last_lba - 33
    fs_sectors    = last_part_lba - PART_LBA + 1

    # GPT header at LBA 1
    hdr = b'EFI PART'                                # signature
    hdr += u32(0x00010000)                            # revision 1.0
    hdr += u32(92)                                    # header size
    hdr += u32(0)  # CRC placeholder                  # header CRC (filled later)
    hdr += u32(0)                                      # reserved
    hdr += u64(1)                                      # my LBA
    hdr += u64(last_lba)                               # backup LBA
    hdr += u64(PART_LBA)                               # first usable LBA
    hdr += u64(last_part_lba)                          # last usable LBA
    hdr += disk_guid                                   # disk GUID
    hdr += u64(2)                                      # partition entry start LBA
    hdr += u32(128)                                    # number of partition entries
    hdr += u32(128)                                    # size of each entry
    hdr += u32(0)  # CRC placeholder                   # partition entries CRC (filled later)
    hdr  = hdr.ljust(92, b'\x00')
    img[GPT_OFF:GPT_OFF+92] = hdr

    # Partition entry: EFI System
    pe  = efi_type
    pe += part_guid
    pe += u64(PART_LBA)
    pe += u64(last_part_lba)
    pe += u64(0)                                       # attributes
    name = "EFI System\0".encode('utf-16-le')
    pe  += name.ljust(72, b'\x00')
    img[GPT_PARTS_OFF:GPT_PARTS_OFF+128] = pe

    # Patch CRCs
    parts_data = bytes(img[GPT_PARTS_OFF:GPT_PARTS_OFF + 128 * 128])
    parts_crc  = crc32(parts_data)
    img[GPT_OFF+88:GPT_OFF+92] = parts_crc
    hdr_data = bytes(img[GPT_OFF:GPT_OFF+92])
    hdr_crc  = crc32(hdr_data[:16] + b'\x00\x00\x00\x00' + hdr_data[20:])
    img[GPT_OFF+16:GPT_OFF+20] = hdr_crc

    # Backup GPT at the last LBA
    backup_lba = last_lba
    img[backup_lba * SECTOR : backup_lba * SECTOR + 92] = img[GPT_OFF:GPT_OFF+92]
    img[backup_lba * SECTOR + 24: backup_lba * SECTOR + 32] = u64(backup_lba)
    # Partition entries backup at last_lba - 32 to last_lba - 1
    img[(last_lba - 32) * SECTOR : (last_lba - 32) * SECTOR + 128*128] = parts_data

    # Write the filesystem into the partition area
    import subprocess, tempfile
    fs_fd, fs_path = tempfile.mkstemp()
    os.close(fs_fd)
    try:
        fs_sectors = last_part_lba - PART_LBA + 1
        with open(fs_path, 'wb') as f:
            f.seek(fs_sectors * SECTOR - 1)
            f.write(b'\x00')
        subprocess.run(['mformat', '-F', '-i', fs_path], check=True, capture_output=True)
        subprocess.run(['mmd', '-i', fs_path, '::/EFI'], check=True, capture_output=True)
        subprocess.run(['mmd', '-i', fs_path, '::/EFI/BOOT'], check=True, capture_output=True)
        subprocess.run(['mcopy', '-i', fs_path, efi_path, '::/EFI/BOOT/'], check=True, capture_output=True)
        with open(fs_path, 'rb') as f:
            fs_data = f.read()
    finally:
        os.unlink(fs_path)

    # Copy filesystem into the partition area
    fs_start = PART_LBA * SECTOR
    img[fs_start:fs_start + len(fs_data)] = fs_data

    with open(out_path, 'wb') as f:
        f.write(img)
    print(f"Wrote {out_path} ({mb} MB)", file=sys.stderr)

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("Usage: mkusb.py bootx64.efi usb.img [size_mb]", file=sys.stderr)
        sys.exit(1)
    build(sys.argv[1], sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else 64)
