#!/bin/sh
# build-uboot-payload.sh — build the krane U-Boot payload for mmcblk0p1:
#
#   0x0000  64-byte arm64 Image header (code0 = b +0x40, image_size,
#           flags bit3, magic at 0x38 — booting.rst contract, verified
#           against depthcharge src/arch/arm/boot64.c)
#   0x0040  uboot-wrapper.S (entry shim: revive display, branch to
#           U-Boot; branch imm26 patched)
#   0x0040+ uboot.bin (u-boot-nodtb.bin + embedded control DTB),
#           contiguous — NO interior padding.
#
# U-Boot's PIE fixup requires runtime _start == link _start (mod 4K);
# CONFIG_TEXT_BASE is chosen so link _start is 4K-aligned (with the
# vectors/stub bytes preceding it), and the file is placed so the
# runtime _start lands on a 4K boundary. Verified via u-boot.sym and
# the ELF section table.
#
# then pack with mkdepthcharge (devkeys) and verify.
set -e
cd "$(dirname "$0")"

UBOOT_BIN="${UBOOT_BIN:-/home/vhaudiquet/u-boot/u-boot.bin}"
UBOOT_SYM="${UBOOT_SYM:-/home/vhaudiquet/u-boot/u-boot.sym}"
DTB="${DTB:-krane-sku176.dtb}"
WRAP_IMG=krane-uboot.bin
OUT_PAYLOAD=krane-uboot-payload.bin

# 1. assemble the diagnostic wrapper (raw binary, no relocations)
aarch64-linux-gnu-gcc -c uboot-wrapper.S -o uboot-wrapper.o
aarch64-linux-gnu-objcopy -O binary uboot-wrapper.o uboot-wrapper.bin
WRAP_LEN=$(stat -c %s uboot-wrapper.bin)

# 2. header + wrapper + u-boot, patch the wrapper's `b .` branch
python3 - "$UBOOT_BIN" "$UBOOT_SYM" "$WRAP_IMG" "$WRAP_LEN" <<'EOF'
import struct, subprocess, sys

uboot_path, sym_path, out = sys.argv[1], sys.argv[2], sys.argv[3]
wrap_len = int(sys.argv[4])
uboot = open(uboot_path, 'rb').read()
wrapper = open('uboot-wrapper.bin', 'rb').read()
assert len(wrapper) == wrap_len

# u-boot.bin starts at __image_copy_start (the lowest output VMA);
# _start's file offset is its delta from that base. The PIE fixup in
# start.S loads the link base from _TEXT_BASE and the run base from
# adr _start, so _start MUST equal __image_copy_start (a 4-byte
# linker fill sneaks in when CONFIG_TEXT_BASE is not 8-aligned —
# start.o's .text input section is 8-aligned — and then every
# relocated pointer is skewed by 4). Fail loudly instead.
copy_vma = None
start_vma = None
for line in open(sym_path):
    parts = line.split()
    if len(parts) >= 2 and parts[-1] == '__image_copy_start':
        copy_vma = int(parts[0], 16)
    elif len(parts) >= 2 and parts[-1] == '_start':
        start_vma = int(parts[0], 16)
if copy_vma is None or start_vma is None:
    raise SystemExit('symbols not found in u-boot.sym')
assert start_vma == copy_vma, \
    ("_start 0x%x != __image_copy_start 0x%x: CONFIG_TEXT_BASE is not "
     "8-aligned and a linker fill shifted _start" %
     (start_vma, copy_vma))
assert start_vma % 0x1000 == 0, \
    "link _start 0x%x not 4K-aligned" % start_vma
start_file_off = start_vma - copy_vma
wrap_off = 0x40
uboot_off = 0x1000
while (uboot_off + start_file_off) % 0x1000:
    uboot_off += 0x1000
pad = uboot_off - wrap_off - wrap_len
assert pad >= 0
total = uboot_off + len(uboot)
assert (0x40000000 + uboot_off + start_file_off) % 0x1000 == 0
hdr = bytearray(64)
hdr[0:4] = struct.pack('<I', (0x40 >> 2) | 0x14000000)  # code0: b +0x40
struct.pack_into('<Q', hdr, 0x10, total)                # image_size
struct.pack_into('<Q', hdr, 0x18, 1 << 3)               # flags: bit3
hdr[0x38:0x3c] = b'ARM\x64'                             # magic

# patch the wrapper's `b .` — searched, NOT assumed to be the last
# instruction (v3 has the delay subroutine after it; patching blind
# overwrites delay's `ret` and the first blink never returns).
hits = [i for i in range(0, len(wrapper), 4)
        if struct.unpack_from('<I', wrapper, i)[0] == 0x14000000]
assert len(hits) == 1, \
    "expected exactly one `b .` in wrapper, found %d" % len(hits)
br_off = wrap_off + hits[0]
imm = (uboot_off + start_file_off - br_off) // 4
wrapper = bytearray(wrapper)
wrapper[hits[0]:hits[0] + 4] = \
    struct.pack('<I', (imm & 0x03ffffff) | 0x14000000)
with open(out, 'wb') as f:
    f.write(hdr)
    f.write(wrapper)
    f.write(bytes(pad))
    f.write(uboot)
print("layout: header 64, wrapper %d (0x40..0x%x), uboot %d @0x%x, "
      "total %d, runtime _start 0x%x" %
      (wrap_len, uboot_off, len(uboot), uboot_off, total,
       0x40000000 + uboot_off + start_file_off))
EOF

# 3. pack + verify
PYTHONPATH=/root/krane-fb-stub/src/depthcharge-tools python3 -m depthcharge_tools.mkdepthcharge \
    -A arm64 \
    -o "$OUT_PAYLOAD" \
    -n "krane u-boot framebuffer console" \
    -d "$WRAP_IMG" \
    -b "$DTB"

echo "---- verify ----"
futility vbutil_kernel --verify "$OUT_PAYLOAD"
sha256sum "$OUT_PAYLOAD"