"""Make a Zig macOS dylib loadable as an iOS arm64 dylib."""
import struct
from pathlib import Path

PAGE = 0x4000
MH_MAGIC_64 = 0xFEEDFACF
LC_SEGMENT_64 = 0x19
LC_LOAD_DYLINKER = 0xE
LC_BUILD_VERSION = 0x32


def patch_dylib(path: Path) -> None:
    data = bytearray(path.read_bytes())
    if data[:4] != struct.pack("<I", MH_MAGIC_64):
        raise SystemExit(f"{path} is not a 64-bit Mach-O")
    ncmds, sizeofcmds = struct.unpack_from("<II", data, 16)
    off = 32
    cmds = []
    for _ in range(ncmds):
        cmd, cmdsize = struct.unpack_from("<II", data, off)
        cmds.append((cmd, bytearray(data[off : off + cmdsize])))
        off += cmdsize

    segments = []
    for cmd, chunk in cmds:
        if cmd != LC_SEGMENT_64:
            continue
        name = bytes(chunk[8:24]).split(b"\x00", 1)[0].decode()
        vmaddr, vmsize, _fileoff, filesize = struct.unpack_from("<QQQQ", chunk, 24)
        segments.append((name, vmaddr, vmsize, filesize, chunk))

    for index, (name, vmaddr, vmsize, filesize, chunk) in enumerate(segments):
        if name == "__LINKEDIT":
            aligned = 0x3000
            while aligned < filesize:
                aligned += PAGE
        else:
            aligned = (vmsize + PAGE - 1) & ~(PAGE - 1)
            next_vm = segments[index + 1][1]
            if vmaddr + aligned > next_vm:
                raise SystemExit(f"{name} cannot be aligned without overlapping the next segment")
        if aligned < filesize:
            raise SystemExit(f"{name} aligned size {aligned:#x} is smaller than file size {filesize:#x}")
        struct.pack_into("<Q", chunk, 32, aligned)
        print(f"{name} vmsize {vmsize:#x} -> {aligned:#x}")

    rebuilt = bytearray()
    removed_dylinker = False
    patched_platform = False
    for cmd, chunk in cmds:
        if cmd == LC_BUILD_VERSION:
            struct.pack_into("<III", chunk, 8, 2, 0x000F0000, 0x000F0000)
            patched_platform = True
        if cmd == LC_LOAD_DYLINKER:
            removed_dylinker = True
            print("removed LC_LOAD_DYLINKER")
            continue
        rebuilt += chunk
    if not removed_dylinker:
        raise SystemExit("LC_LOAD_DYLINKER was not found; the dylib layout changed")
    if not patched_platform:
        raise SystemExit("LC_BUILD_VERSION was not found")

    data[32 : 32 + len(rebuilt)] = rebuilt
    data[32 + len(rebuilt) : 32 + sizeofcmds] = b"\x00" * (sizeofcmds - len(rebuilt))
    struct.pack_into("<II", data, 16, ncmds - 1, len(rebuilt))
    path.write_bytes(data)