#!/usr/bin/env bash
# Install the MCom03R ALT P11 demonstration rootfs on a GPT/UEFI disk.
set -Eeuo pipefail

readonly ESP_GUID=C12A7328-F81F-11D2-BA4B-00A0C93EC93B
ESP_SIZE_MIB=${ESP_SIZE_MIB:-300}
TITLE=${TITLE:-ALT P11 MCom03R GPU/VPU PoC}
TIMEOUT=${TIMEOUT:-10}
KERNEL=${KERNEL:-Image-6.6.30-buildroot}
FALLBACK_KERNEL=${FALLBACK_KERNEL:-vmlinuz-6.6.107-mcom03r-6.6-alt1}
FALLBACK_INITRD=${FALLBACK_INITRD:-initrd-6.6.107-mcom03r-6.6-alt1.img}
PRIMARY_DTB=${PRIMARY_DTB:-/boot/elvees-6.6.30-buildroot/mcom03r-elvmc03smarc-r3.0.0-elvsmarccb-r3.3.0.dtb}
FALLBACK_DTB=${FALLBACK_DTB:-/usr/lib/devicetree/6.6.107-mcom03r-6.6-alt1/mcom03r-elvmc03smarc-r3.0.0-elvsmarccb-r3.3.0.dtb}
HOSTNAME_VALUE=
TAR_VERBOSE=()

log_info() { printf '[INFO] %s\n' "$*"; }
log_success() { printf '[SUCCESS] %s\n' "$*"; }
die() { printf '[ERROR] %s\n' "$*" >&2; exit 1; }

usage() {
    cat <<'EOF'
Usage: tar2dev_mcom03r_uefi.sh [options] ROOTFS.tar.xz /dev/DISK

Destructively create a GPT disk with an EFI System Partition and the ALT P11
MCom03R demonstration rootfs. The default boot entry uses the validated
Buildroot 6.6.30 kernel without an initrd. The original ALT kernel is retained
as a fallback GRUB entry when its files exist in the rootfs.

Options:
  -h            Show this help.
  -v            List files while extracting the rootfs.
  -n HOSTNAME   Set /etc/hostname in the installed rootfs.

Environment overrides:
  KERNEL, INITRD, PRIMARY_EXTRA_ARGS, FALLBACK_KERNEL, FALLBACK_INITRD,
  PRIMARY_DTB, FALLBACK_DTB, TITLE, TIMEOUT, ESP_SIZE_MIB.
EOF
}

if [[ ${1:-} == --help ]]; then
    usage
    exit 0
fi

while getopts ':hvn:' option; do
    case $option in
        h) usage; exit 0 ;;
        v) TAR_VERBOSE=(-v) ;;
        n) HOSTNAME_VALUE=$OPTARG ;;
        :) die "Option -$OPTARG requires a value." ;;
        *) die "Unknown option: -$OPTARG" ;;
    esac
done
shift $((OPTIND - 1))

[[ $# -eq 2 ]] || { usage >&2; exit 64; }
ROOTFS_ARCHIVE=$1
DEVICE=$2

[[ $EUID -eq 0 ]] || die "Run this command as root."
[[ -r $ROOTFS_ARCHIVE ]] || die "Unreadable rootfs archive: $ROOTFS_ARCHIVE"
[[ -b $DEVICE ]] || die "Not a block device: $DEVICE"
[[ -e /sys/class/block/${DEVICE##*/} ]] || die "Unknown block device: $DEVICE"
[[ ! -e /sys/class/block/${DEVICE##*/}/partition ]] || \
    die "Pass a whole disk, not a partition: $DEVICE"
[[ -w $DEVICE ]] || die "Block device is not writable: $DEVICE"

for command in awk blkid cp find lsblk mkfs.ext4 mkfs.vfat mount sfdisk tar umount; do
    command -v "$command" >/dev/null 2>&1 || die "Required command is absent: $command"
done

partition_path() {
    local disk=$1 number=$2 separator=
    [[ $disk =~ [0-9]$ ]] && separator=p
    printf '%s%s%s\n' "$disk" "$separator" "$number"
}

boot_name() {
    local name=$1
    name=${name#../}
    name=${name#/}
    name=${name#boot/}
    printf '%s\n' "$name"
}

unmount_device() {
    local mountpoint
    while IFS= read -r mountpoint; do
        [[ -n $mountpoint ]] || continue
        log_info "Unmounting $mountpoint..."
        umount "$mountpoint"
    done < <(lsblk -lnpo MOUNTPOINT "$DEVICE" | \
        awk 'NF { path[NR]=$0 } END { for (i=NR; i>0; i--) print path[i] }')
}

ROOT_DIR=
ESP_DIR=
ROOT_MOUNTED=0
ESP_MOUNTED=0
cleanup() {
    set +e
    ((ESP_MOUNTED)) && umount "$ESP_DIR"
    ((ROOT_MOUNTED)) && umount "$ROOT_DIR"
    [[ -z $ESP_DIR ]] || rmdir "$ESP_DIR" 2>/dev/null
    [[ -z $ROOT_DIR ]] || rmdir "$ROOT_DIR" 2>/dev/null
}
trap cleanup EXIT INT TERM

log_info "Unmounting partitions on $DEVICE..."
unmount_device

log_info "Partitioning $DEVICE as GPT..."
printf 'label: gpt\n,%sM,%s,*\n,,,\n' "$ESP_SIZE_MIB" "$ESP_GUID" | \
    sfdisk -f "$DEVICE" >/dev/null
command -v partprobe >/dev/null 2>&1 && partprobe "$DEVICE" || true
command -v udevadm >/dev/null 2>&1 && udevadm settle || true

ESP_PART=$(partition_path "$DEVICE" 1)
ROOT_PART=$(partition_path "$DEVICE" 2)
for _ in {1..50}; do
    [[ -b $ESP_PART && -b $ROOT_PART ]] && break
    sleep 0.1
done
[[ -b $ESP_PART && -b $ROOT_PART ]] || die "New partitions did not appear."

log_info "Formatting $ESP_PART as FAT32..."
mkfs.vfat -F32 -n ESP "$ESP_PART" >/dev/null
log_info "Formatting $ROOT_PART as ext4..."
mkfs.ext4 -F -L root "$ROOT_PART" >/dev/null

ROOT_UUID=$(blkid -s UUID -o value "$ROOT_PART")
ROOT_PARTUUID=$(sfdisk --part-uuid "$DEVICE" 2)
ROOT_PARTUUID_LOWER=${ROOT_PARTUUID,,}
log_info "Root UUID: $ROOT_UUID"
log_info "Root PARTUUID: $ROOT_PARTUUID"

ROOT_DIR=$(mktemp -d)
ESP_DIR=$(mktemp -d)
mount "$ROOT_PART" "$ROOT_DIR"
ROOT_MOUNTED=1

log_info "Extracting rootfs..."
tar -C "$ROOT_DIR" "${TAR_VERBOSE[@]}" -xf "$ROOTFS_ARCHIVE"

KERNEL_NAME=$(boot_name "$KERNEL")
if [[ ! -v PRIMARY_EXTRA_ARGS ]]; then
    case $KERNEL_NAME in
        Image-6.6.30-buildroot|Image-6.6.30-malidp-smart-cursor)
            PRIMARY_EXTRA_ARGS='mali_dp.smart_cursor=1 mali_dp.smart_cursor_async=1 cma=512M'
            ;;
        *)
            PRIMARY_EXTRA_ARGS=
            ;;
    esac
fi
if [[ ! -v INITRD ]]; then
    case $KERNEL_NAME in
        Image-6.6.30-buildroot) INITRD= ;;
        *) INITRD=initrd.img ;;
    esac
fi
INITRD_NAME=$(boot_name "$INITRD")
FALLBACK_KERNEL_NAME=$(boot_name "$FALLBACK_KERNEL")
FALLBACK_INITRD_NAME=$(boot_name "$FALLBACK_INITRD")

[[ -f $ROOT_DIR/boot/$KERNEL_NAME ]] || \
    die "Selected kernel is absent from rootfs: /boot/$KERNEL_NAME"
if [[ -n $INITRD_NAME && ! -f $ROOT_DIR/boot/$INITRD_NAME ]]; then
    die "Selected initrd is absent from rootfs: /boot/$INITRD_NAME"
fi
[[ $PRIMARY_DTB == /* ]] || die "PRIMARY_DTB must be an absolute rootfs path."
[[ -f $ROOT_DIR$PRIMARY_DTB ]] || \
    die "Selected primary DTB is absent from rootfs: $PRIMARY_DTB"

touch "$ROOT_DIR/boot/$ROOT_PARTUUID.findme"
mount "$ESP_PART" "$ESP_DIR"
ESP_MOUNTED=1
mkdir -p "$ESP_DIR/EFI/BOOT"

GRUB_EFI=$(find "$ROOT_DIR" -type f \
    \( -name grubaa64.efi -o -name grubx64.efi \) -print -quit)
[[ -n $GRUB_EFI ]] || die "No GRUB EFI loader was found in the rootfs."
log_info "Using GRUB EFI loader: ${GRUB_EFI#"$ROOT_DIR"}"
cp "$GRUB_EFI" "$ESP_DIR/EFI/BOOT/bootaa64.efi"

DTB_SOURCE=$ROOT_DIR/boot/elvees-6.6.30-buildroot
if [[ ! -d $DTB_SOURCE ]]; then
    DTB_FILE=$(find "$ROOT_DIR" -type f -name 'mcom03*.dtb' -print -quit)
    [[ -n $DTB_FILE ]] || die "No MCom03 DTBs were found in the rootfs."
    DTB_SOURCE=$(dirname "$DTB_FILE")
fi
mkdir -p "$ESP_DIR/dtb"
cp -a "$DTB_SOURCE"/. "$ESP_DIR/dtb/"
log_info "Copied DTBs from ${DTB_SOURCE#"$ROOT_DIR"}."

GRUB_CFG=$ESP_DIR/EFI/BOOT/grub.cfg
cat >"$GRUB_CFG" <<EOF
set default="0"
set timeout="$TIMEOUT"
set menu_color_normal=white/black
set menu_color_highlight=black/light-gray
terminal_output console
terminal_input console
EOF

write_menu_entry() {
    local title=$1 kernel=$2 initrd=$3 dtb=$4 extra_args=${5:-}
    cat >>"$GRUB_CFG" <<EOF

menuentry "$title" {
  search --file --set=root /boot/$ROOT_PARTUUID.findme
  if [ "\$?" = 0 ]; then
    echo "Loading $title ..."
    devicetree $dtb
    linux /boot/$kernel root=PARTUUID=$ROOT_PARTUUID_LOWER console=ttyS0,115200n8 console=tty0 loglevel=8 ignore_loglevel ro rootwait $extra_args
EOF
    if [[ -n $initrd ]]; then
        printf '    initrd /boot/%s\n' "$initrd" >>"$GRUB_CFG"
    fi
    cat >>"$GRUB_CFG" <<'EOF'
  fi
}
EOF
}

write_menu_entry "$TITLE" "$KERNEL_NAME" "$INITRD_NAME" \
    "$PRIMARY_DTB" "$PRIMARY_EXTRA_ARGS"
if [[ -f $ROOT_DIR/boot/$FALLBACK_KERNEL_NAME &&
        -f $ROOT_DIR/boot/$FALLBACK_INITRD_NAME ]]; then
    [[ $FALLBACK_DTB == /* ]] || \
        die "FALLBACK_DTB must be an absolute rootfs path."
    [[ -f $ROOT_DIR$FALLBACK_DTB ]] || \
        die "Selected fallback DTB is absent from rootfs: $FALLBACK_DTB"
    write_menu_entry "ALT P11 fallback kernel" \
        "$FALLBACK_KERNEL_NAME" "$FALLBACK_INITRD_NAME" "$FALLBACK_DTB" \
        "systemd.unit=multi-user.target"
fi

mkdir -p "$ROOT_DIR/etc"
cat >"$ROOT_DIR/etc/fstab" <<EOF
proc        /proc     proc    nosuid,noexec,gid=proc                        0 0
devpts      /dev/pts  devpts  nosuid,noexec,gid=tty,mode=620,ptmxmode=0666  0 0
tmpfs       /tmp      tmpfs   nosuid                                        0 0
UUID=$ROOT_UUID /     ext4    defaults                                      1 1
EOF
if [[ -n $HOSTNAME_VALUE ]]; then
    printf '%s\n' "$HOSTNAME_VALUE" >"$ROOT_DIR/etc/hostname"
fi

sync
log_success "Disk is ready. Primary kernel: $KERNEL_NAME; initrd: ${INITRD_NAME:-none}"
log_info "Generated EFI configuration: /EFI/BOOT/grub.cfg"