#!/usr/bin/env bash
# mkdiskimage.sh -- build a single-partition FAT32 boot disk image.
#
# Used by `make boot_image TARGET=<board>` (kernel/Makefile). One script for
# every board so the geometry rules live in one place:
#   - partition 1 starts at 1 MiB (sector 2048) and runs to the last usable
#     sector; the FAT32 filesystem is sized to the partition exactly (a
#     larger filesystem overruns the partition -- the bug that made early
#     SER5 thumbdrive images unbootable, see kernel/Makefile `thumbdrive`);
#   - GPT tables reserve 33 sectors at the end of the disk for the backup
#     header, MBR tables do not.
#
# Usage:
#   mkdiskimage.sh --out IMG --table gpt|mbr --size-mib N --label LABEL \
#                  SRC:DEST [SRC:DEST ...]
#
# --table gpt  -> partition type "EFI System"     (UEFI firmware, U-Boot bootefi)
#
# The StarshipOS stick and data stick (docs/v4.0.0/STICK.md 4; --table gpt only):
#   --starship IDENT_MIB --mkartemis TOOL --esp-mib E
# lays out an ESP of E MiB holding the SRC:DEST files (E=0: no ESP, no files),
# then a StarshipOS identity partition of IDENT_MIB MiB (all zeros), then a
# StarshipOS Artemis partition to the last whole MiB before the backup GPT,
# formatted and signed by TOOL (tools/mkartemis.c). Every partition starts and
# ends on a whole MiB, so on 1 KiB blocks, as the kernel's GPT reader needs.
# --table mbr  -> partition type 0x0c FAT32 (LBA)  (Raspberry Pi firmware)
# DEST is an absolute path inside the FAT filesystem; parent directories are
# created as needed.
#
# Needs only sfdisk (util-linux), mkfs.fat (dosfstools) and mtools -- no
# root, no loop devices.
set -euo pipefail

OUT="" TABLE="" SIZE_MIB="" LABEL="" IDENT_MIB="" MKARTEMIS="" ESP_MIB=""
while [ $# -gt 0 ]; do
    case "$1" in
        --out)      OUT="$2"; shift 2 ;;
        --table)    TABLE="$2"; shift 2 ;;
        --size-mib) SIZE_MIB="$2"; shift 2 ;;
        --label)    LABEL="$2"; shift 2 ;;
        --starship) IDENT_MIB="$2"; shift 2 ;;
        --mkartemis) MKARTEMIS="$2"; shift 2 ;;
        --esp-mib)  ESP_MIB="$2"; shift 2 ;;
        --) shift; break ;;
        -*) echo "mkdiskimage.sh: unknown option $1" >&2; exit 2 ;;
        *) break ;;
    esac
done

[ -n "$OUT" ] && [ -n "$TABLE" ] && [ -n "$SIZE_MIB" ] && [ -n "$LABEL" ] || {
    echo "mkdiskimage.sh: --out, --table, --size-mib and --label are required" >&2; exit 2; }
if [ -n "$IDENT_MIB" ]; then
    [ "$TABLE" = gpt ] && [ -n "$MKARTEMIS" ] && [ -n "$ESP_MIB" ] || {
        echo "mkdiskimage.sh: --starship needs --table gpt, --mkartemis and --esp-mib" >&2; exit 2; }
    [ "$ESP_MIB" -gt 0 ] || [ $# -eq 0 ] || { echo "mkdiskimage.sh: --esp-mib 0 takes no files" >&2; exit 2; }
fi
[ $# -gt 0 ] || [ "${ESP_MIB:-}" = 0 ] || { echo "mkdiskimage.sh: no SRC:DEST files given" >&2; exit 2; }

for tool in sfdisk mkfs.fat mmd mcopy; do
    command -v "$tool" >/dev/null 2>&1 || PATH="$PATH:/usr/sbin:/sbin"
    command -v "$tool" >/dev/null 2>&1 || {
        echo "mkdiskimage.sh: $tool not found (install util-linux, dosfstools, mtools)" >&2; exit 1; }
done

TOTAL_SECTORS=$(( SIZE_MIB * 2048 ))
START=2048
case "$TABLE" in
    gpt) LAST=$(( TOTAL_SECTORS - 34 )); PTYPE="C12A7328-F81F-11D2-BA4B-00A0C93EC93B" ;;
    mbr) LAST=$(( TOTAL_SECTORS - 1 ));  PTYPE="c" ;;
    *) echo "mkdiskimage.sh: --table must be gpt or mbr" >&2; exit 2 ;;
esac
PART_SECTORS=$(( LAST - START + 1 ))
# StarshipOS partition types (kernel/src/gpt.c)
IDENT_TYPE="23BEBFE1-4817-4CF3-A537-9069BF159788"
ART_TYPE="BCD7B586-1C8E-4EEF-86A3-273A8D8EE624"
if [ -n "$IDENT_MIB" ]; then
    PART_SECTORS=$(( ESP_MIB * 2048 ))
    IDENT_START=$(( START + PART_SECTORS ))
    IDENT_SECTORS=$(( IDENT_MIB * 2048 ))
    ART_START=$(( IDENT_START + IDENT_SECTORS ))
    ART_SECTORS=$(( (LAST + 1 - ART_START) / 2048 * 2048 ))
    [ "$ART_SECTORS" -gt 0 ] || { echo "mkdiskimage.sh: no room for Artemis in ${SIZE_MIB} MiB" >&2; exit 2; }
fi

ESP="$OUT.part1"
rm -f "$OUT" "$ESP"
mkdir -p "$(dirname "$OUT")"

truncate -s $(( TOTAL_SECTORS * 512 )) "$OUT"
if [ -n "$IDENT_MIB" ]; then
    {
        echo "label: gpt"
        [ "$ESP_MIB" -gt 0 ] && printf 'start=%d, size=%d, type=%s, name="EFI System"\n' "$START" "$PART_SECTORS" "$PTYPE"
        printf 'start=%d, size=%d, type=%s, name="StarshipOS identity"\n' "$IDENT_START" "$IDENT_SECTORS" "$IDENT_TYPE"
        printf 'start=%d, size=%d, type=%s, name="StarshipOS Artemis"\n' "$ART_START" "$ART_SECTORS" "$ART_TYPE"
    } | sfdisk --quiet "$OUT"
    "$MKARTEMIS" --format "$OUT" --first-lba "$ART_START" --sectors "$ART_SECTORS" >/dev/null || {
        echo "mkdiskimage.sh: formatting Artemis failed" >&2; rm -f "$OUT"; exit 1; }
    if [ "$ESP_MIB" = 0 ]; then
        echo "  IMAGE   $OUT (gpt, ${SIZE_MIB} MiB: StarshipOS identity ${IDENT_MIB} MiB, Artemis $(( ART_SECTORS / 2048 )) MiB)"
        exit 0
    fi
elif [ "$TABLE" = gpt ]; then
    printf 'label: gpt\nstart=%d, size=%d, type=%s, name="EFI System"\n' \
        "$START" "$PART_SECTORS" "$PTYPE" | sfdisk --quiet "$OUT"
else
    printf 'label: dos\nstart=%d, size=%d, type=%s, bootable\n' \
        "$START" "$PART_SECTORS" "$PTYPE" | sfdisk --quiet "$OUT"
fi

truncate -s $(( PART_SECTORS * 512 )) "$ESP"
mkfs.fat -F 32 -n "$LABEL" "$ESP" >/dev/null

for spec in "$@"; do
    src="${spec%%:*}"
    dest="${spec#*:}"
    [ -f "$src" ] || { echo "mkdiskimage.sh: missing input $src" >&2; rm -f "$OUT" "$ESP"; exit 1; }
    dir="$(dirname "$dest")"
    path=""
    IFS=/ read -r -a parts <<< "${dir#/}"
    for p in "${parts[@]}"; do
        [ -n "$p" ] || continue
        path="$path/$p"
        mmd -D s -i "$ESP" "::$path" 2>/dev/null || true
    done
    mcopy -o -i "$ESP" "$src" "::$dest"
done

dd if="$ESP" of="$OUT" bs=512 seek="$START" conv=notrunc status=none
rm -f "$ESP"
if [ -n "$IDENT_MIB" ]; then
    echo "  IMAGE   $OUT (gpt, ${SIZE_MIB} MiB: ESP ${ESP_MIB} MiB FAT32 '$LABEL', StarshipOS identity ${IDENT_MIB} MiB, Artemis $(( ART_SECTORS / 2048 )) MiB)"
else
    echo "  IMAGE   $OUT ($TABLE, ${SIZE_MIB} MiB, FAT32 '$LABEL')"
fi