#!/usr/bin/env bash

set -euo pipefail


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

CONFIG_FILE="NenjimHub.conf"

JAVA_ARCHIVE="zulu25.36.205-ca-jre25.0.4.1-linux_musl_aarch64.tar.gz"
JAVA_DIRECTORY="zulu25.36.205-ca-jre25.0.4.1-linux_musl_aarch64"
JAVA_VERSION_DIRECTORY="jre-25.0.4.1"
JAVA_MAJOR_DIRECTORY="jre-25"

ZREPL_VERSION="v0.7.0"
ZREPL_ASSET="zrepl-linux-arm64"
ZREPL_LOCAL_BINARY="$ZREPL_ASSET"

# Cards marketed as 16 GB are commonly a little smaller than 16,000,000,000
# usable bytes. 14 GiB actual capacity is the baseline; the configured
# partition and zvol sizes are also checked below.
MINIMUM_DEVICE_BYTES=$((14 * 1024 * 1024 * 1024))

# Raspberry Pi Zero 2 W boots most reliably from a classic MBR partition table.
# 0x0c is FAT32 LBA, 0x82 is Linux swap, and 0xbf identifies Solaris/ZFS.
BOOT_PARTITION_TYPE="c"
SWAP_PARTITION_TYPE="82"
ZFS_PARTITION_TYPE="bf"

TEMP_DIRECTORY=""
STAGING_ROOT=""
VIRT_MOUNT=""
POOL_CREATED=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 "$VIRT_MOUNT" ]] && mountpoint -q "$VIRT_MOUNT" 2>/dev/null; then
        umount "$VIRT_MOUNT"
    fi

    if (( POOL_CREATED == 1 )) && zpool list -H -o name "$POOL_NAME" >/dev/null 2>&1; then
        zpool export "$POOL_NAME"
    fi

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

download_zrepl()
{
    local destination="$1"
    local metadata asset_data asset_url asset_digest actual_digest

    echo "Downloading official zrepl ${ZREPL_VERSION} ARM64 binary"

    metadata="$({
        curl \
            --fail \
            --location \
            --proto '=https' \
            --silent \
            --show-error \
            "https://api.github.com/repos/zrepl/zrepl/releases/tags/${ZREPL_VERSION}"
    })"

    asset_data="$({
        printf '%s' "$metadata" |
            ZREPL_ASSET="$ZREPL_ASSET" python3 -c '
import json
import os
import sys

release = json.load(sys.stdin)
name = os.environ["ZREPL_ASSET"]

for asset in release.get("assets", []):
    if asset.get("name") == name:
        print(asset.get("browser_download_url", ""))
        print(asset.get("digest", ""))
        break
else:
    raise SystemExit(f"Release asset not found: {name}")
'
    })" || die "Could not read the zrepl release metadata."

    asset_url="$(printf '%s\n' "$asset_data" | sed -n '1p')"
    asset_digest="$(printf '%s\n' "$asset_data" | sed -n '2p')"

    [[ "$asset_url" == https://github.com/zrepl/zrepl/releases/download/* ]] ||
        die "Unexpected zrepl download URL returned by GitHub."

    [[ "$asset_digest" =~ ^sha256:[[:xdigit:]]{64}$ ]] ||
        die "GitHub did not provide a SHA-256 digest for $ZREPL_ASSET."

    curl \
        --fail \
        --location \
        --proto '=https' \
        --silent \
        --show-error \
        --output "$destination" \
        "$asset_url"

    actual_digest="sha256:$(sha256sum "$destination" | awk '{print $1}')"

    [[ "$actual_digest" == "$asset_digest" ]] ||
        die "SHA-256 validation failed for the downloaded zrepl binary."
}


# ==============================================================================
# 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 command in \
    awk blkid blockdev curl file findmnt install lsblk mkfs.ext4 mkfs.vfat mkswap \
    mount mountpoint partprobe python3 sed sfdisk sha256sum tar udevadm \
    umount wipefs zfs zpool
do
    command -v "$command" >/dev/null 2>&1 ||
        die "Missing required command: $command"
done

[[ -f "$JAVA_ARCHIVE" ]] || die "Missing input file: $JAVA_ARCHIVE"
[[ -b "$DEVICE" ]] || die "Not a block device: $DEVICE"

for name in \
    BOOT_PARTITION_SIZE_MIB VIRT_ZVOL_SIZE_MIB \
    SWAP_SIZE_MIB_PART SWAP_SIZE_MIB_ZFS SWAP_SIZE_PCT_ZRAM
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."
(( 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"
[[ "$BOOT_LABEL" =~ ^[[:alnum:]_-]{1,11}$ ]] ||
    die "BOOT_LABEL must be 1-11 letters, digits, underscores or hyphens."
[[ "$SWAP_PART_LABEL" =~ ^[[:alnum:]_-]{1,15}$ ]] ||
    die "SWAP_PART_LABEL must be 1-15 letters, digits, underscores or hyphens."

DEVICE_BYTES="$(blockdev --getsize64 "$DEVICE")"

CONFIGURED_MINIMUM_DEVICE_BYTES=$((
    (1 + BOOT_PARTITION_SIZE_MIB + SWAP_SIZE_MIB_PART +
        VIRT_ZVOL_SIZE_MIB + SWAP_SIZE_MIB_ZFS + 2048) * 1024 * 1024
))

if (( CONFIGURED_MINIMUM_DEVICE_BYTES > MINIMUM_DEVICE_BYTES )); then
    REQUIRED_DEVICE_BYTES="$CONFIGURED_MINIMUM_DEVICE_BYTES"
else
    REQUIRED_DEVICE_BYTES="$MINIMUM_DEVICE_BYTES"
fi

(( DEVICE_BYTES >= REQUIRED_DEVICE_BYTES )) ||
    die "The card has $DEVICE_BYTES bytes; the selected layout needs at least $REQUIRED_DEVICE_BYTES bytes."

if lsblk -nrpo MOUNTPOINT "$DEVICE" | grep -q '[^[:space:]]'; then
    die "A partition on $DEVICE is mounted. Unmount it before continuing."
fi

if zpool list -H -o name "$POOL_NAME" >/dev/null 2>&1; then
    die "A pool named $POOL_NAME is already imported on this PC."
fi

if zpool status -P 2>/dev/null | grep -Fq "$DEVICE"; then
    die "$DEVICE belongs to a pool that is currently imported. Export it first."
fi

HOST_ZFS_VERSION="$(zpool --version | sed -n 's/^zfs-\([^ -]*\).*/\1/p' | head -1)"

[[ "$HOST_ZFS_VERSION" == "$EXPECTED_ZFS_VERSION" ]] ||
    die "This design expects OpenZFS $EXPECTED_ZFS_VERSION; the PC has ${HOST_ZFS_VERSION:-unknown}."

COMPATIBILITY_FILE="/usr/share/zfs/compatibility.d/$POOL_COMPATIBILITY"
[[ -f "$COMPATIBILITY_FILE" ]] ||
    die "Missing OpenZFS compatibility file: $COMPATIBILITY_FILE"

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

TEMP_DIRECTORY="$(mktemp -d /tmp/nenjim-storage.XXXXXX)"
STAGING_ROOT="$TEMP_DIRECTORY/root"
VIRT_MOUNT="$TEMP_DIRECTORY/virt"

mkdir -p "$STAGING_ROOT" "$VIRT_MOUNT"
trap cleanup EXIT


# ==============================================================================
# Destructive-action confirmation
# ==============================================================================

echo
echo "WARNING: ALL DATA ON ${DEVICE} WILL BE DESTROYED."
echo
echo "The new layout will be:"
echo "  partition 1: ${BOOT_PARTITION_SIZE_MIB} MiB FAT32 Alpine boot partition"
if (( SWAP_SIZE_MIB_PART > 0 )); then
    echo "  partition 2: ${SWAP_SIZE_MIB_PART} MiB raw Linux swap (priority ${SWAP_PRI_PART})"
else
    echo "  partition 2: omitted (raw swap is disabled)"
fi
echo "  partition 3: remaining space, OpenZFS pool $POOL_NAME"
if (( SWAP_SIZE_MIB_ZFS > 0 )); then
    echo "  $POOL_NAME/swap: ${SWAP_SIZE_MIB_ZFS} MiB (priority ${SWAP_PRI_ZFS})"
else
    echo "  $POOL_NAME/swap: disabled"
fi
echo "  ZRAM: ${SWAP_SIZE_PCT_ZRAM}% of RAM (priority ${SWAP_PRI_ZRAM})"
echo
lsblk "$DEVICE"
echo
read -r -p "Type YES to create the new storage layout: " answer

[[ "$answer" == "YES" ]] || die "Aborted."
[[ -b "$DEVICE" ]] || die "The block device disappeared: $DEVICE"


# ==============================================================================
# Create the physical partitions
# ==============================================================================

echo
echo "=== Creating FAT32, optional raw swap and OpenZFS partitions ==="

wipefs --all "$DEVICE"

LOGICAL_SECTOR_BYTES="$(blockdev --getss "$DEVICE")"

case "$LOGICAL_SECTOR_BYTES" in
    512|4096)
        ;;
    *)
        die "Unsupported logical sector size: $LOGICAL_SECTOR_BYTES bytes."
        ;;
esac

BOOT_START_SECTOR=$((1024 * 1024 / LOGICAL_SECTOR_BYTES))
BOOT_SIZE_SECTORS=$((BOOT_PARTITION_SIZE_MIB * 1024 * 1024 / LOGICAL_SECTOR_BYTES))
SWAP_SIZE_SECTORS=$((SWAP_SIZE_MIB_PART * 1024 * 1024 / LOGICAL_SECTOR_BYTES))
NEXT_START_SECTOR=$((BOOT_START_SECTOR + BOOT_SIZE_SECTORS))

if (( SWAP_SIZE_MIB_PART > 0 )); then
    SWAP_START_SECTOR="$NEXT_START_SECTOR"
    ZFS_START_SECTOR=$((SWAP_START_SECTOR + SWAP_SIZE_SECTORS))
else
    ZFS_START_SECTOR="$NEXT_START_SECTOR"
fi

{
    echo 'label: dos'
    echo 'unit: sectors'
    echo
    echo "$BOOT_PARTITION : start=$BOOT_START_SECTOR, size=$BOOT_SIZE_SECTORS, type=$BOOT_PARTITION_TYPE, bootable"
    if (( SWAP_SIZE_MIB_PART > 0 )); then
        echo "$SWAP_PARTITION : start=$SWAP_START_SECTOR, size=$SWAP_SIZE_SECTORS, type=$SWAP_PARTITION_TYPE"
    fi
    echo "$ZFS_PARTITION : start=$ZFS_START_SECTOR, type=$ZFS_PARTITION_TYPE"
} | sfdisk --wipe always "$DEVICE"

partprobe "$DEVICE"
udevadm settle

wait_for_block_device "$BOOT_PARTITION" ||
    die "Boot partition did not appear: $BOOT_PARTITION"

if (( SWAP_SIZE_MIB_PART > 0 )); then
    wait_for_block_device "$SWAP_PARTITION" ||
        die "Swap partition did not appear: $SWAP_PARTITION"
fi

wait_for_block_device "$ZFS_PARTITION" ||
    die "ZFS partition did not appear: $ZFS_PARTITION"

mkfs.vfat -F 32 -n "$BOOT_LABEL" "$BOOT_PARTITION"

if (( SWAP_SIZE_MIB_PART > 0 )); then
    mkswap -L "$SWAP_PART_LABEL" "$SWAP_PARTITION"
fi


# ==============================================================================
# Create the ZFS pool and datasets
# ==============================================================================

echo
echo "=== Creating OpenZFS pool $POOL_NAME ==="

zpool create \
    -f \
    -R "$STAGING_ROOT" \
    -o "compatibility=$POOL_COMPATIBILITY" \
    -o ashift=12 \
    -O xattr=sa \
    -O compression=lz4 \
    -O atime=off \
    -O canmount=off \
    -O mountpoint=none \
    "$POOL_NAME" \
    "$ZFS_PARTITION"

POOL_CREATED=1

zfs create -o mountpoint=/home "$POOL_NAME/home"
zfs create -o mountpoint=/usr/local/software "$POOL_NAME/software"
zfs create -o mountpoint=/media/sysop "$POOL_NAME/sysop"


# ==============================================================================
# Install Java on rpool/software
# ==============================================================================

echo
echo "=== Installing Azul Zulu JRE on rpool/software ==="

SOFTWARE_ROOT="$STAGING_ROOT/usr/local/software"
mkdir -p "$SOFTWARE_ROOT"

tar -xzf "$JAVA_ARCHIVE" -C "$SOFTWARE_ROOT"

[[ -d "$SOFTWARE_ROOT/$JAVA_DIRECTORY" ]] ||
    die "The Java archive did not create $JAVA_DIRECTORY."

ln -s "$JAVA_DIRECTORY" "$SOFTWARE_ROOT/$JAVA_VERSION_DIRECTORY"
ln -s "$JAVA_VERSION_DIRECTORY" "$SOFTWARE_ROOT/$JAVA_MAJOR_DIRECTORY"


# ==============================================================================
# Install the official static zrepl ARM64 binary on rpool/software
# ==============================================================================

echo
echo "=== Installing zrepl ${ZREPL_VERSION} on rpool/software ==="

ZREPL_VERSION_NUMBER="${ZREPL_VERSION#v}"
ZREPL_VERSION_DIRECTORY="$SOFTWARE_ROOT/zrepl/versions/$ZREPL_VERSION_NUMBER"
ZREPL_TEMP_BINARY="$TEMP_DIRECTORY/$ZREPL_ASSET"

mkdir -p "$ZREPL_VERSION_DIRECTORY"

if [[ -f "$ZREPL_LOCAL_BINARY" ]]; then
    echo "Using local zrepl binary: $ZREPL_LOCAL_BINARY"
    cp "$ZREPL_LOCAL_BINARY" "$ZREPL_TEMP_BINARY"
else
    download_zrepl "$ZREPL_TEMP_BINARY"
fi

file "$ZREPL_TEMP_BINARY" | grep -Eq 'ELF 64-bit.*ARM aarch64' ||
    die "$ZREPL_ASSET is not an ARM64 ELF binary."

install -m 0755 "$ZREPL_TEMP_BINARY" "$ZREPL_VERSION_DIRECTORY/zrepl"
ln -s "versions/$ZREPL_VERSION_NUMBER" "$SOFTWARE_ROOT/zrepl/current"


# ==============================================================================
# Create and format the shared Docker/containerd zvol
# ==============================================================================

echo
echo "=== Creating ${POOL_NAME}/virt (${VIRT_ZVOL_SIZE_MIB} MiB, ext4) ==="

zfs create \
    -V "${VIRT_ZVOL_SIZE_MIB}M" \
    -o volmode=dev \
    -o logbias=throughput \
    -o primarycache=metadata \
    -o secondarycache=none \
    "$POOL_NAME/virt"

udevadm settle

VIRT_DEVICE="/dev/zvol/$POOL_NAME/virt"
wait_for_block_device "$VIRT_DEVICE" ||
    die "The zvol device did not appear: $VIRT_DEVICE"

mkfs.ext4 -F -L NENJIMVIRT "$VIRT_DEVICE"
mount -t ext4 -o rw,noatime "$VIRT_DEVICE" "$VIRT_MOUNT"

mkdir -p "$VIRT_MOUNT/docker" "$VIRT_MOUNT/containerd"
chmod 0711 "$VIRT_MOUNT/docker" "$VIRT_MOUNT/containerd"

sync
umount "$VIRT_MOUNT"


# ==============================================================================
# Verification and clean export
# ==============================================================================

echo
echo "=== Verifying and exporting the pool ==="

zpool get -H -o value compatibility "$POOL_NAME" |
    grep -qx "$POOL_COMPATIBILITY" ||
    die "The pool compatibility setting is incorrect."

for dataset in home software sysop virt; do
    zfs list "$POOL_NAME/$dataset" >/dev/null ||
        die "Missing ZFS object: $POOL_NAME/$dataset"
done

JAVA_BINARY="$SOFTWARE_ROOT/$JAVA_DIRECTORY/bin/java"

[[ -x "$JAVA_BINARY" ]] ||
    die "The Java archive did not provide an executable bin/java."

# The preparation PC is normally amd64 and therefore cannot execute this
# Raspberry Pi binary. Verify its format and target architecture instead.
file "$JAVA_BINARY" | grep -Eq 'ELF 64-bit.*ARM aarch64' ||
    die "The installed Java binary is not an ARM64 ELF executable."

zpool set cachefile=none "$POOL_NAME"
zpool export "$POOL_NAME"
POOL_CREATED=0

sync


# ==============================================================================
# Finished
# ==============================================================================

cleanup
trap - EXIT

echo
echo "============================================================================"
echo "Storage preparation is complete."
echo
echo "Partition 1 is ready for flash.sh."
if (( SWAP_SIZE_MIB_PART > 0 )); then
    echo "Partition 2 is raw swap and will not be touched by flash.sh."
else
    echo "Partition 2 is intentionally unused."
fi
echo "Partition 3 contains $POOL_NAME and will not be touched by flash.sh."
if (( SWAP_SIZE_MIB_ZFS > 0 )); then
    echo "The Pi will create ${POOL_NAME}/swap (${SWAP_SIZE_MIB_ZFS} MiB) using its real page size."
fi
echo "============================================================================"