#!/usr/bin/env python3
"""
Wrap a raw MTK preloader (GFH / 'MMM') with the EMMC_BOOT + BRLYT boot-region
header that the BootROM looks for at the start of BOOT1.

A raw preloader written straight to BOOT1 has no 'EMMC_BOOT' signature, so the
BootROM can't find it -> black screen / falls back to BROM. SP Flash Tool adds
this 0x800-byte header automatically; this reproduces it.

Structure (verified against a known-good wrapped preloader):
  0x000  "EMMC_BOOT\0\0\0" + ver(1) + next_off(0x200) + 0xFF pad   (512 B)
  0x200  "BRLYT\0\0\0" + ver(1) + gfh_off(0x800) + total_size +
         "BBBB" + 05 00 01 00 + gfh_off(0x800) + total_size + 1 + 0x00 pad
  0x800  the raw preloader (unchanged)

Usage:
    python3 wrap_preloader.py <raw_preloader.bin> [out.img]
"""
import sys, struct, os

HDR = 0x800


def wrap(raw: bytes) -> bytes:
    total = HDR + len(raw)
    h = bytearray(b"\xFF" * HDR)

    # --- EMMC_BOOT region (0x000) ---
    h[0:12] = b"EMMC_BOOT\x00\x00\x00"
    struct.pack_into("<I", h, 0x0C, 1)          # version
    struct.pack_into("<I", h, 0x10, 0x200)      # offset to BRLYT
    # 0x14..0x200 stays 0xFF (as in known-good image)

    # --- BRLYT region (0x200) ---
    for i in range(0x200, HDR):                 # zero-fill the BRLYT region
        h[i] = 0x00
    h[0x200:0x208] = b"BRLYT\x00\x00\x00"
    struct.pack_into("<I", h, 0x208, 1)         # version
    struct.pack_into("<I", h, 0x20C, HDR)       # GFH (preloader) start = 0x800
    struct.pack_into("<I", h, 0x210, total)     # total size
    h[0x214:0x218] = b"BBBB"                     # layout magic
    h[0x218:0x21C] = b"\x05\x00\x01\x00"         # type/flags (from known image)
    struct.pack_into("<I", h, 0x21C, HDR)       # GFH start = 0x800
    struct.pack_into("<I", h, 0x220, total)     # total size
    struct.pack_into("<I", h, 0x224, 1)

    return bytes(h) + raw


def main():
    if len(sys.argv) < 2:
        print(__doc__); return 2
    src = sys.argv[1]
    raw = open(src, "rb").read()
    if raw[:4] == b"EMMC":
        print("Input already has EMMC_BOOT header — nothing to do."); return 0
    if raw[:4] != b"MMM\x01":
        print(f"WARNING: input doesn't start with 'MMM' GFH magic ({raw[:4]!r}); "
              f"wrapping anyway.")
    out = sys.argv[2] if len(sys.argv) > 2 else \
        os.path.splitext(src)[0] + "_wrapped.img"
    wrapped = wrap(raw)
    open(out, "wb").write(wrapped)
    print(f"raw size    : {len(raw)} ({len(raw):#x})")
    print(f"wrapped size: {len(wrapped)} ({len(wrapped):#x})")
    print(f"  0x000: {wrapped[:9]!r}  (must be b'EMMC_BOOT')")
    print(f"  0x200: {wrapped[0x200:0x205]!r}  (must be b'BRLYT')")
    print(f"  0x800: {wrapped[0x800:0x804]!r}  (must be b'MMM\\x01')")
    print(f"-> {out}")
    return 0


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