#!/usr/bin/env bash
# Partition a whole disk with a bringup-oriented Android GBL GPT layout.
#
# Usage:
#   sudo ./partition_gbl_disk.sh [--efi /path/to/BOOTAA64.EFI] /dev/sdX
#
# Optional environment overrides:
#   SUPER_SIZE=8G ./partition_gbl_disk.sh /dev/sdX
#
# The script intentionally asks for confirmation once before touching the disk.

set -euo pipefail

ESP_SIZE="${ESP_SIZE:-8M}"
MISC_SIZE="${MISC_SIZE:-4M}"
BOOT_SIZE="${BOOT_SIZE:-128M}"
VENDOR_BOOT_SIZE="${VENDOR_BOOT_SIZE:-128M}"
INIT_BOOT_SIZE="${INIT_BOOT_SIZE:-8M}"
VBMETA_SIZE="${VBMETA_SIZE:-4M}"
DTBO_SIZE="${DTBO_SIZE:-32M}"
SUPER_SIZE="${SUPER_SIZE:-8G}"
METADATA_SIZE="${METADATA_SIZE:-64M}"

die() {
  echo "error: $*" >&2
  exit 1
}

usage() {
  cat >&2 <<EOF
Usage:
  sudo $0 [--efi /path/to/BOOTAA64.EFI] /dev/sdX

Creates a destructive GPT layout for Android GBL bringup:
  android_esp_a/b, misc, boot_a/b, vendor_boot_a/b, init_boot_a/b,
  vbmeta_a/b, dtbo_a/b, dtb_a/b, pvmfw_a/b, super, metadata, userdata.

Options:
  --efi FILE      Copy FILE to both ESPs as /EFI/BOOT/BOOTAA64.EFI
  -h, --help      Show this help

Environment overrides:
  ESP_SIZE=$ESP_SIZE
  MISC_SIZE=$MISC_SIZE
  BOOT_SIZE=$BOOT_SIZE
  VENDOR_BOOT_SIZE=$VENDOR_BOOT_SIZE
  INIT_BOOT_SIZE=$INIT_BOOT_SIZE
  VBMETA_SIZE=$VBMETA_SIZE
  DTBO_SIZE=$DTBO_SIZE
  PVMFW_SIZE=$PVMFW_SIZE
  SUPER_SIZE=$SUPER_SIZE
  METADATA_SIZE=$METADATA_SIZE
EOF
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || die "missing required command: $1"
}

cleanup() {
  if [[ -n "${ESP_A_MNT:-}" && -d "${ESP_A_MNT:-}" ]]; then
    umount "$ESP_A_MNT" 2>/dev/null || true
    rmdir "$ESP_A_MNT" 2>/dev/null || true
  fi
  if [[ -n "${ESP_B_MNT:-}" && -d "${ESP_B_MNT:-}" ]]; then
    umount "$ESP_B_MNT" 2>/dev/null || true
    rmdir "$ESP_B_MNT" 2>/dev/null || true
  fi
  if [[ -n "${SFDISK_SCRIPT:-}" ]]; then
    rm -f "$SFDISK_SCRIPT"
  fi
}

partition_path() {
  local disk="$1"
  local number="$2"

  case "$disk" in
    *[0-9]) printf '%sp%s\n' "$disk" "$number" ;;
    *) printf '%s%s\n' "$disk" "$number" ;;
  esac
}

EFI_PAYLOAD=""
ARGS=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    --efi)
      [[ $# -ge 2 ]] || die "--efi requires a file path"
      EFI_PAYLOAD="$2"
      shift 2
      ;;
    -h | --help)
      usage
      exit 0
      ;;
    --)
      shift
      ARGS+=("$@")
      break
      ;;
    -*)
      die "unknown argument: $1"
      ;;
    *)
      ARGS+=("$1")
      shift
      ;;
  esac
done

[[ ${#ARGS[@]} -eq 1 ]] || {
  usage
  exit 2
}

DISK="${ARGS[0]}"

[[ $EUID -eq 0 ]] || die "run as root, for example: sudo $0 $DISK"
[[ -b "$DISK" ]] || die "$DISK is not a block device"
if [[ -n "$EFI_PAYLOAD" ]]; then
  [[ -f "$EFI_PAYLOAD" ]] || die "--efi payload does not exist or is not a regular file: $EFI_PAYLOAD"
fi

require_cmd sfdisk
require_cmd lsblk
require_cmd findmnt
if [[ -n "$EFI_PAYLOAD" ]]; then
  require_cmd mount
  require_cmd umount
  require_cmd mkfs.vfat
fi

if [[ "$(lsblk -dn -o TYPE "$DISK")" != "disk" ]]; then
  die "$DISK is not a whole disk according to lsblk"
fi

if findmnt --source "$DISK" >/dev/null 2>&1; then
  die "$DISK itself is mounted; unmount it first"
fi

while read -r child; do
  [[ -n "$child" ]] || continue
  if findmnt --source "$child" >/dev/null 2>&1; then
    die "$child is mounted; unmount all partitions on $DISK first"
  fi
done < <(lsblk -nrpo NAME "$DISK" | tail -n +2)

cat <<EOF
About to DESTROY ALL DATA on:
  $DISK

Planned GPT layout:
  1   android_esp_a    $ESP_SIZE      EFI System, FAT32
  2   android_esp_b    $ESP_SIZE      EFI System, FAT32
  3   misc             $MISC_SIZE
  4   boot_a           $BOOT_SIZE
  5   boot_b           $BOOT_SIZE
  6   vendor_boot_a    $VENDOR_BOOT_SIZE
  7   vendor_boot_b    $VENDOR_BOOT_SIZE
  8   init_boot_a      $INIT_BOOT_SIZE
  9   init_boot_b      $INIT_BOOT_SIZE
  10  vbmeta_a         $VBMETA_SIZE
  11  vbmeta_b         $VBMETA_SIZE
  12  dtbo_a           $DTBO_SIZE
  13  dtbo_b           $DTBO_SIZE
  14  super            $SUPER_SIZE
  15  metadata         $METADATA_SIZE
  16  userdata         remaining space

EOF
if [[ -n "$EFI_PAYLOAD" ]]; then
  cat <<EOF
GBL EFI payload:
  $EFI_PAYLOAD

The payload will be copied to both Android ESPs as:
  /EFI/BOOT/BOOTAA64.EFI

EOF
fi
cat <<EOF
This cannot be undone by this script.
EOF

read -r -p "Continue and destroy all data on $DISK? [y/N]: " CONFIRMATION
case "$CONFIRMATION" in
  y | Y | yes | YES)
    ;;
  *)
    die "aborting"
    ;;
esac

SFDISK_SCRIPT="$(mktemp)"
trap cleanup EXIT

# Declarative GPT layout for sfdisk. Omitted start sectors let sfdisk choose
# aligned starts. Omitted userdata size consumes the remaining disk.
cat >"$SFDISK_SCRIPT" <<EOF
label: gpt

name="android_esp_a", size=$ESP_SIZE,type=uefi, name="EFI"
name="android_esp_b", size=$ESP_SIZE,type=uefi, name="EFI"
name="misc", size=$MISC_SIZE
name="boot_a", size=$BOOT_SIZE
name="boot_b", size=$BOOT_SIZE
name="vendor_boot_a", size=$VENDOR_BOOT_SIZE
name="vendor_boot_b", size=$VENDOR_BOOT_SIZE
name="init_boot_a", size=$INIT_BOOT_SIZE
name="init_boot_b", size=$INIT_BOOT_SIZE
name="vbmeta_a", size=$VBMETA_SIZE
name="vbmeta_b", size=$VBMETA_SIZE
name="dtbo_a", size=$DTBO_SIZE
name="dtbo_b", size=$DTBO_SIZE
name="super", size=$SUPER_SIZE
name="metadata", size=$METADATA_SIZE
name="userdata"
EOF

echo "Writing GPT with sfdisk..."
sfdisk "$DISK" <"$SFDISK_SCRIPT"

if command -v partprobe >/dev/null 2>&1; then
  partprobe "$DISK" || true
fi
if command -v udevadm >/dev/null 2>&1; then
  udevadm settle || true
fi
blockdev --rereadpt "$DISK" 2>/dev/null || true

ESP_A="$(partition_path "$DISK" 1)"
ESP_B="$(partition_path "$DISK" 2)"

if command -v mkfs.fat >/dev/null 2>&1; then
  echo "Formatting Android ESP partitions as FAT32..."
  mkfs.vfat -n EFI "$ESP_A"
  mkfs.vfat -n EFI "$ESP_B"
else
  echo "warning: mkfs.fat not found; format $ESP_A and $ESP_B as FAT32 manually" >&2
fi
 
if [[ -n "$EFI_PAYLOAD" ]]; then
  echo "Copying GBL EFI payload to Android ESP partitions..."
  ESP_A_MNT="$(mktemp -d /tmp/gbl-esp-a.XXXXXX)"
  ESP_B_MNT="$(mktemp -d /tmp/gbl-esp-b.XXXXXX)"

  mount "$ESP_A" "$ESP_A_MNT"
  mount "$ESP_B" "$ESP_B_MNT"

  mkdir -p "$ESP_A_MNT/EFI/BOOT"
  mkdir -p "$ESP_B_MNT/EFI/BOOT"
  cp "$EFI_PAYLOAD" "$ESP_A_MNT/EFI/BOOT/BOOTAA64.EFI"
  cp "$EFI_PAYLOAD" "$ESP_B_MNT/EFI/BOOT/BOOTAA64.EFI"
  sync

  umount "$ESP_A_MNT"
  umount "$ESP_B_MNT"
  rmdir "$ESP_A_MNT" "$ESP_B_MNT"
  ESP_A_MNT=""
  ESP_B_MNT=""
fi

echo
echo "Done. Resulting partition table:"
sfdisk -d "$DISK"
echo
if [[ -n "$EFI_PAYLOAD" ]]; then
  echo "GBL EFI payload copied to both ESPs as /EFI/BOOT/BOOTAA64.EFI."
else
  echo "Copy GBL to both ESPs as /EFI/BOOT/BOOTAA64.EFI before booting."
fi