#!/usr/bin/env bash

set -euo pipefail


# ==============================================================================
# Configuration
# ==============================================================================

VERSION="3.24.1"
BRANCH="3.24"
CONFIG_FILE="NenjimHub.conf"

IMAGE="alpine-rpi-${VERSION}-aarch64.img"
HEADLESS_OVERLAY="headless.apkovl.tar.gz"
WPA_CONFIG="wpa_supplicant.conf"
UNATTENDED_SCRIPT="unattended.sh"
ZEROTIER_INSTALLER="installZeroTier.sh"
BUILD_TIME_FILE="nenjim-build-time"

BUILDER_IMAGE="alpine:${BRANCH}"
KERNEL_MKSQUASHFS_MEMORY="128M"

BOOT_MOUNT="/mnt/usb"

TEMP_DIRECTORY=""
SOURCE_MOUNT=""
LOOP_DEVICE=""
ZEROTIER_PREBUILT_AVAILABLE=0


# ==============================================================================
# Helpers
# ==============================================================================

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

require_uint()
{
    local name="$1"
    local value="${!name-}"

    [[ "$value" =~ ^[0-9]+$ ]] ||
        die "$name must be a non-negative integer."
}

require_swap_priority()
{
    local name="$1"
    local value="${!name-}"

    require_uint "$name"
    (( value <= 32767 )) ||
        die "$name must be between 0 and 32767."
}

partition_path()
{
    local number="$1"

    if [[ "$DEVICE" =~ [0-9]$ ]]; then
        printf '%sp%s\n' "$DEVICE" "$number"
    else
        printf '%s%s\n' "$DEVICE" "$number"
    fi
}

wait_for_block_device()
{
    local path="$1"
    local attempt

    for attempt in {1..50}; do
        [[ -b "$path" ]] && return 0
        sleep 0.2
    done

    return 1
}

cleanup()
{
    set +e

    if [[ -n "$SOURCE_MOUNT" ]] && mountpoint -q "$SOURCE_MOUNT" 2>/dev/null; then
        umount "$SOURCE_MOUNT"
    fi

    if mountpoint -q "$BOOT_MOUNT" 2>/dev/null; then
        umount "$BOOT_MOUNT"
    fi

    if [[ -n "$LOOP_DEVICE" ]]; then
        losetup -d "$LOOP_DEVICE" 2>/dev/null || true
    fi

    if [[ -n "$TEMP_DIRECTORY" && -d "$TEMP_DIRECTORY" ]]; then
        case "$TEMP_DIRECTORY" in
            /tmp/nenjim-flash.*)
                rm -rf -- "$TEMP_DIRECTORY"
                ;;
            *)
                echo "WARNING: Refusing to remove unexpected path: $TEMP_DIRECTORY" >&2
                ;;
        esac
    fi
}

preserve_sysop_home()
{
    local overlay overlay_count extract_root
    local -a final_overlays=()

    mapfile -t final_overlays < <(
        find "$BOOT_MOUNT" \
            -maxdepth 1 \
            -type f \
            -name '*.apkovl.tar.gz' \
            ! -name 'headless.apkovl.tar.gz' \
            -print
    )

    overlay_count="${#final_overlays[@]}"

    if (( overlay_count == 0 )); then
        echo "No final APKoVL found; this appears to be a fresh installation."
        return 0
    fi

    if (( overlay_count > 1 )); then
        printf 'Found more than one final APKoVL:\n' >&2
        printf '  %s\n' "${final_overlays[@]}" >&2
        die "Keep only the active final APKoVL before updating partition 1."
    fi

    overlay="${final_overlays[0]}"
    extract_root="$TEMP_DIRECTORY/old-overlay"
    mkdir -p "$extract_root"

    echo "Preserving /sysop from $(basename "$overlay")"

    tar \
        --extract \
        --gzip \
        --file "$overlay" \
        --directory "$extract_root" \
        --no-same-owner

    if [[ ! -d "$extract_root/sysop" ]]; then
        echo "The previous APKoVL does not contain /sysop; nothing to preserve."
        return 0
    fi

    tar \
        --create \
        --gzip \
        --file "$TEMP_DIRECTORY/sysop-preserve.tar.gz" \
        --directory "$extract_root" \
        sysop
}

run_kernel_builder()
{
    local runtime="$1"
    local kernel_work="$TEMP_DIRECTORY/kernel-work"

    mkdir -p "$kernel_work"

    "$runtime" run \
        --rm \
        --volume "$BOOT_MOUNT:/media/boot" \
        --volume "$kernel_work:/work" \
        --env KERNEL_MKSQUASHFS_MEMORY="$KERNEL_MKSQUASHFS_MEMORY" \
        "$BUILDER_IMAGE" \
        /bin/sh -euxc '
            apk add --no-cache alpine-conf alpine-keys squashfs-tools xz

            # update-kernel creates a temporary aarch64 APK root. The
            # architecture-specific keys are required when the builder
            # container itself is x86_64. The minimal Docker image can have
            # the shared key files removed even though alpine-keys is marked
            # as installed, so restore all package-owned key files first.
            apk fix --no-cache alpine-keys

            APK_KEYS_DIR=/usr/share/apk/keys
            find "$APK_KEYS_DIR" -maxdepth 1 -type f -name "*.rsa.pub" -print
            test -n "$(find "$APK_KEYS_DIR" -maxdepth 1 -type f -name "*.rsa.pub" -print -quit)"
            apk \
                --arch aarch64 \
                --keys-dir "$APK_KEYS_DIR" \
                update

            MKSQUASHFS_OPTS="-processors 1 -mem $KERNEL_MKSQUASHFS_MEMORY" \
            TMPDIR=/work \
                update-kernel \
                    -a aarch64 \
                    -f rpi \
                    -p zfs-rpi \
                    --keys-dir "$APK_KEYS_DIR" \
                    --repositories-file /etc/apk/repositories \
                    -M /media/boot

            test -s /media/boot/boot/vmlinuz-rpi
            test -s /media/boot/boot/initramfs-rpi
            test -s /media/boot/boot/modloop-rpi

            unsquashfs -ll /media/boot/boot/modloop-rpi |
                grep -E "/zfs[.]ko([.](gz|xz|zst))?$"
        '
}


# ==============================================================================
# Initialization and validation
# ==============================================================================

cd -- "$(dirname -- "${BASH_SOURCE[0]}")"

[[ -f "$CONFIG_FILE" ]] || die "Missing shared configuration: $CONFIG_FILE"
# shellcheck source=NenjimHub.conf
source "./$CONFIG_FILE"

[[ $EUID -eq 0 ]] || die "This script must be run as root."

for name in \
    BOOT_PARTITION_SIZE_MIB VIRT_ZVOL_SIZE_MIB \
    SWAP_SIZE_MIB_PART SWAP_SIZE_MIB_ZFS SWAP_SIZE_PCT_ZRAM \
    DISKLESS_ROOT_SIZE_MIB ZFS_ARC_MAX_MIB ZEROTIER_REQUIRED_SWAP_MIB \
    ZEROTIER_MAKE_JOBS ZEROTIER_ENABLE_SSO
do
    require_uint "$name"
done

for name in SWAP_PRI_ZRAM SWAP_PRI_ZFS SWAP_PRI_PART; do
    require_swap_priority "$name"
done

(( BOOT_PARTITION_SIZE_MIB > 0 )) ||
    die "BOOT_PARTITION_SIZE_MIB must be greater than zero."
(( VIRT_ZVOL_SIZE_MIB > 0 )) ||
    die "VIRT_ZVOL_SIZE_MIB must be greater than zero."
(( DISKLESS_ROOT_SIZE_MIB > 0 )) ||
    die "DISKLESS_ROOT_SIZE_MIB must be greater than zero."
(( ZFS_ARC_MAX_MIB > 0 )) ||
    die "ZFS_ARC_MAX_MIB must be greater than zero."
(( ZEROTIER_REQUIRED_SWAP_MIB > 0 )) ||
    die "ZEROTIER_REQUIRED_SWAP_MIB must be greater than zero."
(( ZEROTIER_MAKE_JOBS > 0 )) ||
    die "ZEROTIER_MAKE_JOBS must be greater than zero."
(( ZEROTIER_ENABLE_SSO == 0 || ZEROTIER_ENABLE_SSO == 1 )) ||
    die "ZEROTIER_ENABLE_SSO must be either 0 or 1."
(( SWAP_SIZE_PCT_ZRAM <= 100 )) ||
    die "SWAP_SIZE_PCT_ZRAM must be between 0 and 100."

[[ "$NENJIMHUB_HOSTNAME" =~ ^[A-Za-z0-9]([A-Za-z0-9-]{0,61}[A-Za-z0-9])?$ ]] ||
    die "NENJIMHUB_HOSTNAME must be one 1-63 character DNS hostname label."
[[ "$CONSOLE_KEYMAP_LAYOUT" =~ ^[A-Za-z0-9][A-Za-z0-9_-]*$ ]] ||
    die "Invalid CONSOLE_KEYMAP_LAYOUT: $CONSOLE_KEYMAP_LAYOUT"
[[ "$CONSOLE_KEYMAP_VARIANT" =~ ^[A-Za-z0-9][A-Za-z0-9_-]*$ ]] ||
    die "Invalid CONSOLE_KEYMAP_VARIANT: $CONSOLE_KEYMAP_VARIANT"
[[ "$ZEROTIER_VERSION" =~ ^[0-9]+([.][0-9]+){2}([._-][A-Za-z0-9]+)*$ ]] ||
    die "Invalid ZEROTIER_VERSION: $ZEROTIER_VERSION"
[[ "$ZEROTIER_PREBUILT_BINARY" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] ||
    die "ZEROTIER_PREBUILT_BINARY must be a plain filename without directories."

DISKLESS_ROOT_SIZE="${DISKLESS_ROOT_SIZE_MIB}M"

for command in \
    blkid blockdev cp date find install losetup lsblk mkfs.vfat mount mountpoint \
    sed sync tar udevadm umount
do
    command -v "$command" >/dev/null 2>&1 ||
        die "Missing required command: $command"
done

for file in \
    "$CONFIG_FILE" \
    "$IMAGE" \
    "$HEADLESS_OVERLAY" \
    "$WPA_CONFIG" \
    "$UNATTENDED_SCRIPT" \
    "$ZEROTIER_INSTALLER"
do
    [[ -f "$file" ]] || die "Missing input file: $file"
done

if [[ -f "$ZEROTIER_PREBUILT_BINARY" ]]; then
    [[ -s "$ZEROTIER_PREBUILT_BINARY" ]] ||
        die "Optional ZeroTier binary is empty: $ZEROTIER_PREBUILT_BINARY"
    ZEROTIER_PREBUILT_AVAILABLE=1
    echo "Using optional precompiled ZeroTier binary: $ZEROTIER_PREBUILT_BINARY"
elif [[ -e "$ZEROTIER_PREBUILT_BINARY" ]]; then
    die "Optional ZeroTier binary is not a regular file: $ZEROTIER_PREBUILT_BINARY"
else
    echo "No precompiled ZeroTier binary found; the Pi will compile it on first boot."
fi

if command -v podman >/dev/null 2>&1; then
    CONTAINER_RUNTIME="podman"
elif command -v docker >/dev/null 2>&1; then
    CONTAINER_RUNTIME="docker"
else
    die "Podman or Docker is required to build the matched ARM64 kernel/modloop."
fi

[[ -b "$DEVICE" ]] || die "Not a block device: $DEVICE"

[[ "$(lsblk -dnro PTTYPE "$DEVICE" 2>/dev/null || true)" == "dos" ]] ||
    die "The card does not use the Raspberry Pi-compatible MBR layout. Run the updated prepare-storage.sh first."

if mountpoint -q "$BOOT_MOUNT" 2>/dev/null; then
    die "$BOOT_MOUNT is already in use. Unmount it before continuing."
fi

if lsblk -nrpo MOUNTPOINT "$DEVICE" | grep -q '[^[:space:]]'; then
    die "A partition on $DEVICE is mounted. Export rpool and unmount the card first."
fi

BOOT_PARTITION="$(partition_path 1)"
SWAP_PARTITION="$(partition_path 2)"
ZFS_PARTITION="$(partition_path 3)"

[[ -b "$BOOT_PARTITION" ]] || die "Missing partition 1: run prepare-storage.sh first."
[[ -b "$ZFS_PARTITION" ]] || die "Missing ZFS partition 3: run the updated prepare-storage.sh first."

EXPECTED_BOOT_PARTITION_BYTES=$((BOOT_PARTITION_SIZE_MIB * 1024 * 1024))
ACTUAL_BOOT_PARTITION_BYTES="$(blockdev --getsize64 "$BOOT_PARTITION")"
[[ "$ACTUAL_BOOT_PARTITION_BYTES" == "$EXPECTED_BOOT_PARTITION_BYTES" ]] ||
    die "$BOOT_PARTITION is $ACTUAL_BOOT_PARTITION_BYTES bytes, but $CONFIG_FILE specifies $EXPECTED_BOOT_PARTITION_BYTES bytes."

[[ "$(blkid -s TYPE -o value "$BOOT_PARTITION" 2>/dev/null || true)" == "vfat" ]] ||
    die "$BOOT_PARTITION is not the expected FAT32 Alpine partition."

if (( SWAP_SIZE_MIB_PART > 0 )); then
    [[ -b "$SWAP_PARTITION" ]] ||
        die "Missing configured raw swap partition 2: run prepare-storage.sh first."
    EXPECTED_SWAP_PARTITION_BYTES=$((SWAP_SIZE_MIB_PART * 1024 * 1024))
    ACTUAL_SWAP_PARTITION_BYTES="$(blockdev --getsize64 "$SWAP_PARTITION")"
    [[ "$ACTUAL_SWAP_PARTITION_BYTES" == "$EXPECTED_SWAP_PARTITION_BYTES" ]] ||
        die "$SWAP_PARTITION is $ACTUAL_SWAP_PARTITION_BYTES bytes, but $CONFIG_FILE specifies $EXPECTED_SWAP_PARTITION_BYTES bytes."
    [[ "$(blkid -s TYPE -o value "$SWAP_PARTITION" 2>/dev/null || true)" == "swap" ]] ||
        die "$SWAP_PARTITION is not a Linux swap area."
    [[ "$(blkid -s LABEL -o value "$SWAP_PARTITION" 2>/dev/null || true)" == "$SWAP_PART_LABEL" ]] ||
        die "$SWAP_PARTITION does not have the expected label $SWAP_PART_LABEL."
elif [[ -b "$SWAP_PARTITION" ]]; then
    die "Partition 2 exists although SWAP_SIZE_MIB_PART=0; rerun prepare-storage.sh with the selected layout."
fi

[[ "$(blkid -s TYPE -o value "$ZFS_PARTITION" 2>/dev/null || true)" == "zfs_member" ]] ||
    die "$ZFS_PARTITION is not an OpenZFS pool member."

TEMP_DIRECTORY="$(mktemp -d /tmp/nenjim-flash.XXXXXX)"
SOURCE_MOUNT="$TEMP_DIRECTORY/image-boot"
mkdir -p "$SOURCE_MOUNT" "$BOOT_MOUNT"

trap cleanup EXIT


# ==============================================================================
# Confirmation: only partition 1 is replaced
# ==============================================================================

echo
echo "Partition 1 (${BOOT_PARTITION}) will be replaced with Alpine ${VERSION}."
if (( SWAP_SIZE_MIB_PART > 0 )); then
    echo "Partition 2 (${SWAP_PARTITION}) is raw swap and will NOT be changed."
else
    echo "Partition 2 is unused."
fi
echo "Partition 3 (${ZFS_PARTITION}) and $POOL_NAME will NOT be changed."
echo
lsblk -f "$DEVICE"
echo
read -r -p "Type NENJIMHUB to update partition 1: " answer

[[ "$answer" == "NENJIMHUB" ]] || die "Aborted."


# ==============================================================================
# Attach the Alpine image read-only
# ==============================================================================

echo
echo "=== Reading the Alpine image ==="

LOOP_DEVICE="$(losetup --find --show --partscan --read-only "$IMAGE")"
SOURCE_PARTITION="${LOOP_DEVICE}p1"

wait_for_block_device "$SOURCE_PARTITION" ||
    die "The image partition did not appear: $SOURCE_PARTITION"

mount -o ro "$SOURCE_PARTITION" "$SOURCE_MOUNT"

[[ -f "$SOURCE_MOUNT/cmdline.txt" ]] ||
    die "The image does not contain a Raspberry Pi cmdline.txt."


# ==============================================================================
# Preserve the rescue home from the previous Alpine overlay
# ==============================================================================

echo
echo "=== Preserving the sysop rescue home, if present ==="

mount "$BOOT_PARTITION" "$BOOT_MOUNT"
preserve_sysop_home

ACTUAL_BOOT_LABEL="$(blkid -s LABEL -o value "$BOOT_PARTITION" 2>/dev/null || true)"
BOOT_UUID="$(blkid -s UUID -o value "$BOOT_PARTITION" 2>/dev/null || true)"

[[ "$ACTUAL_BOOT_LABEL" == "$BOOT_LABEL" ]] ||
    die "$BOOT_PARTITION has label '${ACTUAL_BOOT_LABEL:-none}', expected '$BOOT_LABEL'."

umount "$BOOT_MOUNT"


# ==============================================================================
# Replace only the FAT32 Alpine partition
# ==============================================================================

echo
echo "=== Replacing Alpine partition 1 ==="

BOOT_VOLUME_ID="$(printf '%s' "$BOOT_UUID" | tr -d '-')"

if [[ "$BOOT_VOLUME_ID" =~ ^[[:xdigit:]]{8}$ ]]; then
    mkfs.vfat -F 32 -n "$BOOT_LABEL" -i "$BOOT_VOLUME_ID" "$BOOT_PARTITION"
else
    mkfs.vfat -F 32 -n "$BOOT_LABEL" "$BOOT_PARTITION"
fi

mount "$BOOT_PARTITION" "$BOOT_MOUNT"
cp -a "$SOURCE_MOUNT/." "$BOOT_MOUNT/"

CMDLINE_FILE="$BOOT_MOUNT/cmdline.txt"
CMDLINE_WITHOUT_ROOTFLAGS="$(
    sed 's/rootflags=[^[:space:]]*[[:space:]]*//g' "$CMDLINE_FILE"
)"

printf '%s rootflags=size=%s\n' \
    "$CMDLINE_WITHOUT_ROOTFLAGS" \
    "$DISKLESS_ROOT_SIZE" \
    >"$CMDLINE_FILE"


# ==============================================================================
# Build a matched Raspberry Pi kernel/modloop with zfs-rpi on the PC
# ==============================================================================

echo
echo "=== Adding Alpine's matched Raspberry Pi OpenZFS module ==="

if ! run_kernel_builder "$CONTAINER_RUNTIME"; then
    die "Kernel/modloop build failed. The SD card is not ready to boot."
fi


# ==============================================================================
# Install first-boot inputs
# ==============================================================================

echo
echo "=== Installing unattended first-boot files ==="

install -m 0644 \
    "$HEADLESS_OVERLAY" \
    "$BOOT_MOUNT/headless.apkovl.tar.gz"

install -m 0600 \
    "$WPA_CONFIG" \
    "$BOOT_MOUNT/wpa_supplicant.conf"

install -m 0755 \
    "$UNATTENDED_SCRIPT" \
    "$BOOT_MOUNT/unattended.sh"

install -m 0755 \
    "$ZEROTIER_INSTALLER" \
    "$BOOT_MOUNT/installZeroTier.sh"

install -m 0644 \
    "$CONFIG_FILE" \
    "$BOOT_MOUNT/NenjimHub.conf"

if (( ZEROTIER_PREBUILT_AVAILABLE == 1 )); then
    install -m 0755 \
        "$ZEROTIER_PREBUILT_BINARY" \
        "$BOOT_MOUNT/$ZEROTIER_PREBUILT_BINARY"
fi

# Raspberry Pi Zero 2 W has no battery-backed RTC. Give first boot a recent,
# trusted starting point before it contacts HTTPS package repositories.
date -u '+%Y-%m-%d %H:%M:%S' >"$BOOT_MOUNT/$BUILD_TIME_FILE"
chmod 0644 "$BOOT_MOUNT/$BUILD_TIME_FILE"

if [[ -f "$TEMP_DIRECTORY/sysop-preserve.tar.gz" ]]; then
    install -m 0600 \
        "$TEMP_DIRECTORY/sysop-preserve.tar.gz" \
        "$BOOT_MOUNT/sysop-preserve.tar.gz"
fi


# ==============================================================================
# Verification and finish
# ==============================================================================

echo
echo "=== Verifying and syncing partition 1 ==="

for boot_file in \
    bootcode.bin \
    config.txt \
    fixup.dat \
    start.elf \
    boot/vmlinuz-rpi \
    boot/initramfs-rpi \
    boot/modloop-rpi \
    headless.apkovl.tar.gz \
    unattended.sh \
    installZeroTier.sh \
    NenjimHub.conf \
    nenjim-build-time \
    wpa_supplicant.conf
do
    [[ -s "$BOOT_MOUNT/$boot_file" ]] ||
        die "Missing prepared boot file: $boot_file"
done


if (( ZEROTIER_PREBUILT_AVAILABLE == 1 )); then
    [[ -s "$BOOT_MOUNT/$ZEROTIER_PREBUILT_BINARY" ]] ||
        die "Missing copied ZeroTier binary: $ZEROTIER_PREBUILT_BINARY"
fi

sync
umount "$SOURCE_MOUNT"
umount "$BOOT_MOUNT"
losetup -d "$LOOP_DEVICE"
LOOP_DEVICE=""

cleanup
trap - EXIT

echo
echo "============================================================================"
echo "Alpine partition 1 is ready. $POOL_NAME on partition 3 was left untouched."
if (( ZEROTIER_PREBUILT_AVAILABLE == 1 )); then
    echo "ZeroTier $ZEROTIER_VERSION will use the supplied precompiled binary."
else
    echo "ZeroTier $ZEROTIER_VERSION will be compiled by the Pi."
fi
echo
echo "First boot performs the unattended setup and reboots."
echo "After the second boot, use the sysop rescue/admin account:"
echo "  user:     sysop"
echo "  password: sysop"
echo "============================================================================"